Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 model_name = "Qwen/Qwen2.5-0.5B-Instruct"
 
@@ -8,24 +9,36 @@ model = AutoModelForCausalLM.from_pretrained(
     model_name, low_cpu_mem_usage=True, device_map="auto", torch_dtype="auto"
 )
 
-def predict(messages):
+def predict(history):
+    """
+    history: list of [user, bot] message pairs from the Chatbot
+    """
+    # Convert history into the 'messages' format for chat template
+    messages = []
+    for human, bot in history:
+        if human:
+            messages.append({"role": "user", "content": human})
+        if bot:
+            messages.append({"role": "assistant", "content": bot})
+
     text = tokenizer.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
     )
     model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    logger.info(f"Model generation process started at - {process_id}")
     generated_ids = model.generate(**model_inputs, max_new_tokens=512)
     generated_ids = [
-        output_ids[len(input_ids) :]
+        output_ids[len(input_ids):]
         for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
     ]
-    logger.info(f"Model generation process completed [{process_id}]")
 
     reply = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return reply
+    history.append((messages[-1]["content"] if messages else "", reply))
+    return history
 
-iface = gr.Interface(fn=predict, inputs="messages", outputs="reply")
+with gr.Blocks() as server:
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your message here...")
+    msg.submit(predict, [chatbot], chatbot)
 
-# Launch with API access
-iface.launch(server_name="0.0.0.0", server_port=7860, share=False)
+server.launch(server_name="0.0.0.0", server_port=7860, share=False)