Update app.py

app.py

@@ -1,14 +1,16 @@
-import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
-
 from huggingface_hub import login
+
+# Authenticate with Hugging Face token
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
+# Setup environment
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 os.environ['HF_HOME'] = '/tmp/cache'
 
+# Load model and tokenizer
 model_name = "tiiuae/falcon-rw-1b"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -18,6 +20,7 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
+# Create text generation pipeline
 generator = pipeline(
     "text-generation",
     model=model,
@@ -26,10 +29,11 @@ generator = pipeline(
     torch_dtype=torch_dtype
 )
 
+# Main function for generating chat completions
 def generate_chat_completion(message: str, history: list = None):
     """
-    If history is provided as list of [role,content] pairs,
-    reconstruct prompt. Returns updated history list.
+    If history is provided as list of {'role': str, 'content': str} dicts,
+    it reconstructs the full prompt and returns updated history.
     """
     history = history or []
     prompt = ""
@@ -47,18 +51,7 @@ def generate_chat_completion(message: str, history: list = None):
     )
     reply = output[0]['generated_text'].replace(prompt, "").strip()
 
-    # Append to history
+    # Return updated conversation history
     history.append({"role": "user", "content": message})
     history.append({"role": "assistant", "content": reply})
-
     return history
-
-# Adapt Gradio to pass/receive history automatically
-gr.ChatInterface(
-    fn=generate_chat_completion,
-    title="Falcon Chatbot",
-    description="Roleplay-ready chat using Falcon-RW‑1B",
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear"
-).launch()