Update app.py

app.py

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