Update app.py

app.py

@@ -1,35 +1,35 @@
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
-
-# Load Llama-2 model
-model_name = "meta-llama/Llama-2-7b-chat-hf"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
-
-# Define personalities
-personalities = {
-    "Albert Einstein": "You are Albert Einstein, the famous physicist. Speak wisely and humorously.",
-    "Cristiano Ronaldo": "You are Cristiano Ronaldo, the world-famous footballer. You are confident and say ‘Siuuu!’ often.",
-    "Narendra Modi": "You are Narendra Modi, the Prime Minister of India. Speak in a calm, patriotic manner.",
-    "Robert Downey Jr.": "You are Robert Downey Jr., witty, sarcastic, and charismatic."
-}
-
-# Chat function
-def chat(personality, user_input):
-    prompt = f"{personalities[personality]}\nUser: {user_input}\nAI:"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    output = model.generate(**inputs, max_length=200)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
-
-# Gradio UI
-demo = gr.Interface(
-    fn=chat,
-    inputs=["dropdown", "text"],
-    outputs="text",
-    title="Chat with AI Celebrities",
-    description="Select a character and chat with their AI version.",
-    examples=[["Albert Einstein", "What is relativity?"], ["Cristiano Ronaldo", "How do you stay motivated?"]]
-)
-
-demo.launch()
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+
+# Load Mistral-7B model
+model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+
+# Define personalities
+personalities = {
+    "Albert Einstein": "You are Albert Einstein, the famous physicist. Speak wisely and humorously.",
+    "Cristiano Ronaldo": "You are Cristiano Ronaldo, the world-famous footballer. You are confident and say ‘Siuuu!’ often.",
+    "Narendra Modi": "You are Narendra Modi, the Prime Minister of India. Speak in a calm, patriotic manner.",
+    "Robert Downey Jr.": "You are Robert Downey Jr., witty, sarcastic, and charismatic."
+}
+
+# Chat function
+def chat(personality, user_input):
+    prompt = f"{personalities[personality]}\nUser: {user_input}\nAI:"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_length=200)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
+
+# Gradio UI
+demo = gr.Interface(
+    fn=chat,
+    inputs=["dropdown", "text"],
+    outputs="text",
+    title="Drapel – Chat with AI Celebrities",
+    description="Select a character and chat with their AI version.",
+    examples=[["Albert Einstein", "What is relativity?"], ["Cristiano Ronaldo", "How do you stay motivated?"]]
+)
+
+demo.launch()