Update app.py

app.py

@@ -1,33 +1,21 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
 import gradio as gr
+import torch
 
-model_id = "tiiuae/falcon-rw-1b"
+model_id = "tiiuae/falcon-rw-1b"  # small enough to run in Hugging Face Space
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
-# Chat memory
-chat_history = []
-
-# Main chat function
 def chat(user_input, history):
     prompt = ""
-    for i, (user, bot) in enumerate(history):
-        prompt += f"User: {user}\nAssistant: {bot}\n"
-    prompt += f"User: {user_input}\nAssistant:"
-
-    response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
-    answer = response.split("Assistant:")[-1].strip()
-
-    history.append((user_input, answer))
+    for user, bot in history:
+        prompt += f"User: {user}\nBot: {bot}\n"
+    prompt += f"User: {user_input}\nBot:"
+    
+    response = pipe(prompt, max_new_tokens=128, do_sample=True, temperature=0.7)[0]["generated_text"]
+    reply = response.split("Bot:")[-1].strip()
+    history.append((user_input, reply))
     return history, history
 
-# Interface
-chatbot = gr.Chatbot()
-demo = gr.Interface(fn=chat, inputs=["text", "state"], outputs=[chatbot, "state"], 
-                    title="ChatBot", 
-                    description="Normal chatbot like ChatGPT")
-
-if __name__ == "__main__":
-    demo.launch()
+gr.ChatInterface(chat, chatbot=gr.Chatbot(), title="Lightweight Chatbot").launch()