Update app.py

app.py

@@ -1,107 +1,86 @@
 import os
+import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
 
-MODEL_ID = "google/gemma-3-1b-it"
+MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 
-import os
-from huggingface_hub import login
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+torch.set_num_threads(max(1, (os.cpu_count() or 4) - 1))
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-login(token=HF_TOKEN)
+SYSTEM_PROMPT = (
+    "You are a helpful assistant. Answer clearly and concisely. "
+    "If the user asks for tool JSON, return only valid JSON."
+)
 
 print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
+    use_fast=True,
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
 
 print("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,
+    token=HF_TOKEN,
     device_map="cpu",
-    token=HF_TOKEN
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
 )
 
-print("Model loaded!")
-
-SYSTEM_PROMPT = """
-You are an expert documentary writer and cinematic image prompt engineer.
-
-Tasks:
-1. Explain facts in engaging documentary style
-2. Generate cinematic AI image prompts
-3. Create social-media-ready narration
-
-Always:
-- Be descriptive
-- Use vivid imagery
-- Keep responses high quality
-"""
+model.eval()
+print("Model loaded.")
 
+def respond(message, history):
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
 
-def generate(prompt, max_new_tokens, temperature):
-    full_prompt = f"""
-{SYSTEM_PROMPT}
+    for user_msg, assistant_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
 
-User: {prompt}
-Assistant:
-"""
+    messages.append({"role": "user", "content": message})
 
-    inputs = tokenizer(
-        full_prompt,
-        return_tensors="pt"
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
     )
 
-    with torch.no_grad():
+    inputs = tokenizer(prompt, return_tensors="pt")
+
+    with torch.inference_mode():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
+            max_new_tokens=256,
             do_sample=True,
-            top_p=0.95,
-            repetition_penalty=1.1
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.05,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
 
-    response = tokenizer.decode(
-        outputs[0],
-        skip_special_tokens=True
-    )
-
-    response = response.split("Assistant:")[-1].strip()
-
-    return response
-
-
-demo = gr.Interface(
-    fn=generate,
-    inputs=[
-        gr.Textbox(
-            lines=8,
-            label="Prompt",
-            placeholder="Enter your fact or image prompt request..."
-        ),
-        gr.Slider(
-            minimum=64,
-            maximum=1024,
-            value=256,
-            step=32,
-            label="Max New Tokens"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.5,
-            value=0.7,
-            step=0.1,
-            label="Temperature"
-        )
+    new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
+    reply = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+    return reply
+
+demo = gr.ChatInterface(
+    fn=respond,
+    title="Qwen2.5-1.5B CPU Chat",
+    description="Directly loads the model from Hugging Face Hub. No custom model upload needed.",
+    examples=[
+        "Explain black holes in simple words.",
+        "Write a cinematic image prompt for a medieval knight in a storm.",
+        "Set a timer for 10 minutes because pizza is baking.",
     ],
-    outputs=gr.Textbox(
-        lines=20,
-        label="Response"
-    ),
-    title="Gemma 3 4B CPU Demo",
-    description="Running fully on CPU using Hugging Face Spaces"
 )
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+if __name__ == "__main__":
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)