Update app.py

app.py

@@ -1,28 +1,52 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 import os
 
 # Login using Hugging Face token from environment variable (set via Secrets)
 from huggingface_hub import login
-login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
+hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+if hf_token is None:
+    raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable not set")
+login(hf_token)
 
-# Use float32 for CPU/GPU compatibility
-torch_dtype = torch.float32
-
-# Set Hugging Face cache dir (optional but helps in Spaces)
+# Set Hugging Face cache dir (optional but recommended for Spaces)
 os.environ['HF_HOME'] = '/tmp/cache'
 
-# Load Falcon 1B Instruct model pipeline
+model_name = "tiiuae/falcon-rw-1b-instruct"
+
+# Load tokenizer and model explicitly (better control)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Use bfloat16 if on GPU and supported, else fallback to float32
+if torch.cuda.is_available():
+    torch_dtype = torch.bfloat16
+    device_map = "auto"
+else:
+    torch_dtype = torch.float32
+    device_map = None
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch_dtype,
+    device_map=device_map,
+    trust_remote_code=True
+)
+
+# Create text-generation pipeline with tokenizer and model
 generator = pipeline(
     "text-generation",
-    model="tiiuae/falcon-rw-1b-instruct",
-    device=0 if torch.cuda.is_available() else -1,
-    torch_dtype=torch_dtype
+    model=model,
+    tokenizer=tokenizer,
+    device_map=device_map,
+    torch_dtype=torch_dtype,
 )
 
 def generate_chat_completion(message, history):
-    """Simple chatbot logic"""
+    """
+    Simple chat function using Falcon 1B instruct model.
+    Formats prompt for chat style and returns response.
+    """
     prompt = f"User: {message}\nAssistant:"
     output = generator(
         prompt,
@@ -30,17 +54,20 @@ def generate_chat_completion(message, history):
         temperature=0.9,
         top_p=0.9,
         repetition_penalty=1.1,
-        do_sample=True
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
     )
-    response = output[0]['generated_text'].replace(prompt, "").strip()
+    # Remove the prompt from the generated text to get clean assistant reply
+    generated_text = output[0]['generated_text']
+    response = generated_text[len(prompt):].strip()
     return response
 
-# Gradio chat interface
+# Launch Gradio Chat Interface
 gr.ChatInterface(
     fn=generate_chat_completion,
-    title="Falcon Chatbot",
+    title="Falcon 1B Instruct Chatbot",
     description="Roleplay-ready chat using Falcon-1B-Instruct",
     retry_btn="Retry",
     undo_btn="Undo",
-    clear_btn="Clear"
+    clear_btn="Clear",
 ).launch()