Update app.py

app.py

@@ -7,22 +7,35 @@ model_id = "mistralai/Mistral-7B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
 
-def chat(prompt, history=[]):
-    system_prompt = "You are a helpful AI assistant."
-    full_prompt = system_prompt + "\n"
+# Define modes
+modes = {
+    "Standard Mode": {
+        "prompt": "You are a helpful, honest, and friendly assistant.",
+        "temperature": 0.7
+    },
+    "Brainstorming Mode": {
+        "prompt": "You are a creative, out-of-the-box thinker helping users brainstorm ideas, explore thoughts, and be wildly imaginative.",
+        "temperature": 1.0
+    }
+}
+
+def chat(prompt, history=[], mode="Standard Mode"):
+    settings = modes[mode]
+    system_prompt = settings["prompt"]
+    temp = settings["temperature"]
 
+    full_prompt = system_prompt + "\n"
     for user, bot in history:
         full_prompt += f"User: {user}\nAssistant: {bot}\n"
     full_prompt += f"User: {prompt}\nAssistant:"
 
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    streamer = TextStreamer(tokenizer)
 
     output = model.generate(
         **inputs,
         max_new_tokens=300,
         do_sample=True,
-        temperature=0.7,
+        temperature=temp,
         top_p=0.95,
         pad_token_id=tokenizer.eos_token_id
     )
@@ -32,4 +45,18 @@ def chat(prompt, history=[]):
     history.append((prompt, response))
     return response, history
 
-gr.ChatInterface(chat, title="Mistral Chatbot").launch()
+with gr.Blocks() as demo:
+    mode_dropdown = gr.Dropdown(choices=list(modes.keys()), value="Standard Mode", label="Select Mode")
+    chat_interface = gr.ChatInterface(fn=lambda message, history: chat(message, history, mode_dropdown.value))
+
+    mode_dropdown.change(
+        fn=lambda mode: None,
+        inputs=mode_dropdown,
+        outputs=[],
+        queue=False
+    )
+
+    demo.add(mode_dropdown)
+    demo.add(chat_interface)
+
+demo.launch()