Update app.py

app.py

@@ -9,55 +9,71 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MODEL_ID = "le-llm/lapa-v0.1-reasoning-only"
+torch._dynamo.config.disable = True
 
+MODEL_ID = "le-llm/lapa-v0.1-reasoning-only"
 
 
 def load_model():
     """Lazy-load model & tokenizer (for zeroGPU)."""
-    device = "cuda"# if torch.cuda.is_available() else "cpu"
+    device = "cuda"  # if torch.cuda.is_available() else "cpu"
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        #torch_dtype=torch.bfloat16, #if device == "cuda" else torch.float32,
-        device_map="auto"# if device == "cuda" else None,
-    ) #.cuda()
+        torch_dtype=torch.bfloat16,  # if device == "cuda" else torch.float32,
+        device_map="auto",  # if device == "cuda" else None,
+    )  # .cuda()
     print(f"Selected device:", device)
     return model, tokenizer, device
 
+
 # Load model/tokenizer each request → allows zeroGPU to cold start & then release
 model, tokenizer, device = load_model()
 
+
+def user(user_message, history: list):
+    return "", history + [{"role": "user", "content": user_message}]
+
+
+def append_example_message(x: gr.SelectData, history):
+    print(x)
+    print(x.value)
+    print(x.value["text"])
+    if x.value["text"] is not None:
+        history.append({"role": "user", "content": x.value["text"]})
+
+    return history
+
+
 @spaces.GPU
-def respond(
-    message,
+def bot(
     history: list[dict[str, str]],
-    max_tokens,
-    temperature,
-    top_p,
+    # max_tokens,
+    # temperature,
+    # top_p,
 ):
-    
+
     # [{"role": "system", "content": system_message}] +
     # Build conversation
-    messages =  history + [
-        {"role": "user", "content": message}
-    ]
+    max_tokens = 4096
+    temperature = 0.7
+    top_p = 0.95
 
-    input_text = tokenizer.apply_chat_template(
-        messages,
+    input_text: str = tokenizer.apply_chat_template(
+        history,
         tokenize=False,
         add_generation_prompt=True,
-        enable_thinking=True,
+        # enable_thinking=True,
     )
-    input_text += "<think>" # TODO: remove short term fix
-    print(input_text)
-    inputs = tokenizer(input_text, return_tensors="pt").to(model.device) # .to(device)
 
+    input_text = input_text.replace(tokenizer.bos_token, "", 1)
+    print(input_text)
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)  # .to(device)
+    print("Decoded input:", tokenizer.decode(inputs["input_ids"][0]))
+    print([{id: tokenizer.decode([id])} for id in inputs["input_ids"][0]])
     # Streamer setup
     streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_special_tokens=True,
-        skip_prompt=True
+        tokenizer, skip_prompt=True  # skip_special_tokens=True  # ,
     )
 
     # Run model.generate in background thread
@@ -66,25 +82,64 @@ def respond(
         max_new_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
+        top_k=64,
         do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
+        # eos_token_id=tokenizer.eos_token_id,
         streamer=streamer,
     )
     thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
+    history.append({"role": "assistant", "content": ""})
     # Yield tokens as they come in
-    partial_output = ""
     for new_text in streamer:
-        partial_output += new_text
-        yield partial_output
+        history[-1]["content"] += new_text
+        yield history
+
 
+import gradio as gr
+import random
+import time
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(
+        type="messages",
+        allow_tags=["think"],
+        examples=[
+            {"text": i}
+            for i in [
+                "хто тримає цей район?",
+                "Напиши історію про Івасика-Телесика",
+                "Яка найвища гора в Україні?",
+                "Як звали батька Тараса Григоровича Шевченка?",
+                # "Як можна заробити нелегально швидко гроші?"],
+                "Яка з цих гір не знаходиться у Європі? Говерла, Монблан, Гран-Парадізо, Еверест",
+                "Дай відповідь на питання\nЧому у качки жовті ноги?",
+            ]
+        ],
+    )
+    msg = gr.Textbox(label="Message", autofocus=True)
+    send_btn = gr.Button("Send")
+    # clear = gr.Button("Clear")
 
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"),
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=True).then(
+        bot, chatbot, chatbot
+    )
+
+    chatbot.example_select(
+        append_example_message, [chatbot], [chatbot], queue=True
+    ).then(bot, chatbot, chatbot)
+
+    send_btn.click(user, [msg, chatbot], [msg, chatbot], queue=True).then(
+        bot, chatbot, chatbot
+    )
+
+    # clear.click(lambda: None, None, chatbot, queue=True)
+
+if __name__ == "__main__":
+    demo.launch()
+
+"""gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(
             minimum=0.1,
@@ -92,19 +147,4 @@ chatbot = gr.ChatInterface(
             value=0.95,
             step=0.05,
             label="Top-p (nucleus sampling)",
-        ),
-    ],
-    examples=[
-        ["хто тримає цей район?"],
-        ["Напиши історію про Івасика-Телесика"],
-        ["Яка найвища гора в Україні?"],
-        ["Як звали батька Тараса Григоровича Шевченка?"],
-        #["Як можна заробити нелегально швидко гроші?"],
-        ["Яка з цих гір не знаходиться у Європі? Говерла, Монблан, Гран-Парадізо, Еверест"],
-        [
-        "Дай відповідь на питання\nЧому у качки жовті ноги?"
-    ]],
-)
-
-if __name__ == "__main__":
-    chatbot.launch()
+        ),"""