Update app.py

app.py

@@ -1,62 +1,43 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-model_id = "mistralai/Mistral-7B-Instruct"
+model_id = "openchat/openchat-3.5-1210"
 
+# Load model
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
-
-# 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"
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.float16
+)
+
+# System prompt
+def build_prompt(history, user_input):
+    system_prompt = "<|system|>\nYou are a helpful AI assistant.\n"
+    messages = system_prompt
+
     for user, bot in history:
-        full_prompt += f"User: {user}\nAssistant: {bot}\n"
-    full_prompt += f"User: {prompt}\nAssistant:"
+        messages += f"<|user|>\n{user}\n<|assistant|>\n{bot}\n"
+    messages += f"<|user|>\n{user_input}\n<|assistant|>\n"
+    return messages
 
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+def chat(user_input, history=[]):
+    prompt = build_prompt(history, user_input)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     output = model.generate(
         **inputs,
         max_new_tokens=300,
+        temperature=0.7,
         do_sample=True,
-        temperature=temp,
-        top_p=0.95,
+        top_p=0.9,
         pad_token_id=tokenizer.eos_token_id
     )
 
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-    response = decoded.split("Assistant:")[-1].strip()
-    history.append((prompt, response))
-    return response, history
-
-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)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    answer = response.split("<|assistant|>")[-1].strip()
+    history.append((user_input, answer))
+    return answer, history
 
-demo.launch()
+gr.ChatInterface(chat, title="OpenChat AI Assistant").launch()