Update app.py

app.py

@@ -17,7 +17,6 @@ PLACEHOLDER = """
 <p>Hi! I'm MisMath. A Math advisor. My model is based on mathstral-7B-v0.1. Feel free to ask your questions</p>
 <p>Mathstral 7B is a model specializing in mathematical and scientific tasks, based on Mistral 7B.</p>
 <p>mathstral-7B-v0.1 is the first Mathstral model</p>
-<img src="Mistral.png" alt="MathStral Model" style="width:300px;height:200px;">
 </center>
 """
 
@@ -68,15 +67,12 @@ def stream_chat(
     top_k: int = 20,
     penalty: float = 1.2,
 ):
-    print(f'message: {message}')
-    print(f'history: {history}')
-
     # Prepare the conversation context
     conversation_text = system_prompt + "\n"
     for prompt, answer in history:
-        conversation_text += f"User: {prompt}\nAssistant: {answer}\n"
+        conversation_text += f"User: {prompt}\nMisMath: {answer}\n"
 
-    conversation_text += f"User: {message}\nAssistant:"
+    conversation_text += f"User: {message}\nMisMath:"
 
     # Tokenize the conversation text
     input_ids = tokenizer(conversation_text, return_tensors="pt").input_ids.to(model.device)
@@ -102,8 +98,11 @@ def stream_chat(
     for new_text in streamer:
         buffer += new_text
         # Clean the buffer to remove unwanted prefixes
-        cleaned_text = buffer.split("Assistant:")[-1].strip()
-        yield cleaned_text
+        if "MisMath:" in buffer:
+            cleaned_text = buffer.split("MisMath:")[-1].strip()
+            yield cleaned_text
+        else:
+            yield buffer.strip()
 
 # Define the Gradio chatbot component
 chatbot = gr.Chatbot(height=500, placeholder=PLACEHOLDER)