Update app.py

app.py

@@ -18,36 +18,47 @@ def respond(
     temperature,
     top_p,
 ):
-    # ... (rest of the respond function remains the same)
+    history = history[-3:]
+    # Construct the prompt
+    prompt = f"<s>{system_message}\n\n"
+    for user_msg, assistant_msg in history:
+        prompt += f"<|user|>{user_msg}<|end|></s> <|assistant|>{assistant_msg}<|end|></s>"
+    prompt += f"<|user|>{message}<|end|></s> <|assistant|>"
 
+    # Generate response
+    response = ""
+    for token in model(
+        prompt,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        stream=True,
+        stop=["<|end|>", "</s>"]
+    ):
+        response += token['choices'][0]['text']
+        yield response.strip()
+        
 # Create the Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Med TinyLlama Chat")
-    gr.Markdown("Chat with the Med TinyLlama model for medical information.")
-    
-    chatbot = gr.Chatbot(height=300)
-    msg = gr.Textbox(label="Type your message here", placeholder="Ask a medical question...")
-    
-    with gr.Accordion("Advanced Options", open=False):
-        system_message = gr.Textbox(value="", label="System message")
-        max_tokens = gr.Slider(minimum=128, maximum=4096, value=2048, step=1, label="Max new tokens")
-        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
-        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p (nucleus sampling)")
-    
-    submit_btn = gr.Button("Send")
-    clear_btn = gr.Button("Clear")
-
-    submit_btn.click(
-        respond,
-        inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p],
-        outputs=[chatbot],
-    )
-    msg.submit(
-        respond,
-        inputs=[msg, chatbot, system_message, max_tokens, temperature, top_p],
-        outputs=[chatbot],
-    )
-    clear_btn.click(lambda: None, None, chatbot, queue=False)
+demo = gr.ChatInterface(
+    respond,
+    undo_btn="Отменить",
+    clear_btn="Очистить",
+    additional_inputs=[
+        # gr.Textbox(value="You are a friendly medical assistant.", label="System message"),
+        gr.Textbox(value="", label="System message"),
+        gr.Slider(minimum=128, maximum=4096, value=2048, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.9,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+    ],
+    title="Med TinyLlama Chat",
+    description="Chat with the Med TinyLlama model for medical information.",
+)
 
 if __name__ == "__main__":
     demo.launch()