Update app.py

app.py

@@ -26,11 +26,13 @@ def generate(
         tokenizer.eos_token = "<|im_end|>"
         print(tokenizer)
         pipe.tokenizer = tokenizer
-
-        prompt = f"<|im_start|>system\n{system}<|im_end|>\n"
-        for (user_turn, assistant_turn) in history:
-            prompt += f"<|im_start|>user\n{user_turn}<|im_end|>\n<|im_start|>assistant\n{assistant_turn}<|im_end|>\n"
-        prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+        if model_name == "Locutusque/TinyMistral-248M":
+            prompt = prompt
+        else:
+            prompt = f"<|im_start|>system\n{system}<|im_end|>\n"
+            for (user_turn, assistant_turn) in history:
+                prompt += f"<|im_start|>user\n{user_turn}<|im_end|>\n<|im_start|>assistant\n{assistant_turn}<|im_end|>\n"
+            prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
 
         streamer = TextIteratorStreamer(pipe.tokenizer, timeout=240.0, skip_prompt=True, skip_special_tokens=True)
         generation_kwargs = dict(
@@ -60,7 +62,7 @@ def generate(
         print(f"An error occurred: {e}")
         yield "An error occurred during generation."
 
-model_choices = ["Locutusque/Apollo-2.0-Llama-3.1-8B", "Locutusque/Hercules-6.1-StockQwen-2.5-7B", "Locutusque/Hercules-6.1-Llama-3.1-8B", "Locutusque/DareQwen-2.5-7B", "Locutusque/Hercules-5.0-Index-1.9B", "Locutusque/StockQwen-2.5-7B"]
+model_choices = ["Locutusque/Apollo-2.0-Llama-3.1-8B", "Locutusque/TinyMistral-248M", "Locutusque/Hercules-6.1-Llama-3.1-8B", "Locutusque/DareQwen-2.5-7B", "Locutusque/Hercules-5.0-Index-1.9B", "Locutusque/StockQwen-2.5-7B"]
 # What at the best options? 
 g = gr.ChatInterface(
     fn=generate,