Update app.py

app.py

@@ -4,11 +4,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 from huggingface_hub import login
 
-# Load API token securely
+# Load Hugging Face API token securely
 HF_TOKEN = os.getenv("HF_TOKEN")  # Read token from environment variable
 login(token=HF_TOKEN)
 
-# Define personalities BEFORE using them in the dropdown
+# ✅ Using a lightweight Llama-2 model that works on CPU
+model_name = "TheBloke/Llama-2-7B-Chat-GGML"  # 4-bit quantized model (CPU-friendly)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, device_map="cpu")
+
+# 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.",
@@ -16,28 +21,26 @@ personalities = {
     "Robert Downey Jr.": "You are Robert Downey Jr., witty, sarcastic, and charismatic."
 }
 
-# Load Llama-2 Model
-model_name = "meta-llama/Llama-2-7b-chat-hf"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", use_auth_token=True)
-
 # 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)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")  # Running on CPU
+    output = model.generate(**inputs, max_length=300)
+    
+    response_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response_text
 
-# ✅ Ensure the dropdown has predefined choices
+# Gradio UI
 demo = gr.Interface(
     fn=chat,
     inputs=[
         gr.Dropdown(choices=list(personalities.keys()), label="Choose a Celebrity"),
-        "text"
+        gr.Textbox(label="Your Message")
     ],
-    outputs="text",
+    outputs=gr.Textbox(label="AI Response"),
     title="Drapel – Chat with AI Celebrities",
     description="Select a character and chat with their AI version.",
 )
 
+# Launch app
 demo.launch()