Update app.py

app.py

@@ -1,24 +1,17 @@
 import gradio as gr
-import os
-
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 model_id = "tiiuae/falcon-rw-1b"
 
-os.makedirs("offload", exist_ok=True)
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", offload_folder="./offload")
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-def chat_with_expert(message, history):
-    prompt = f"<s>[INST] You are an expert assistant. Answer with clarity and depth.\n{message} [/INST]"
-    response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]['generated_text']
-    answer = response.split('[/INST]')[-1].strip()
-    history.append((message, answer))
+def chat_with_user(message, history=[]):
+    prompt = message
+    result = generator(prompt, max_new_tokens=200)[0]['generated_text']
+    history.append((message, result.strip()))
     return history, history
 
-chatbot = gr.ChatInterface(fn=chat_with_expert, title="Expert Chat Assistant")
-chatbot.launch()
+gr.ChatInterface(chat_with_user, title="Expert Chat Assistant").launch()