Update app.py

app.py

@@ -1,34 +1,69 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-MODEL_ID = "google/flan-t5-base"   # swap to flan-t5-base if you have more CPU
-pipe = pipeline("text2text-generation", model=MODEL_ID)
+# Tiny, modern instruct model that can (patiently) run on CPU
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 
-def infer(prompt, max_new_tokens, num_beams, length_penalty):
-    if not prompt.strip():
+# Load tokenizer + model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float32,      # CPU-safe; on GPU you could use torch.float16/bfloat16
+    low_cpu_mem_usage=True          # helps reduce peak RAM on load
+)
+
+# Make sure a pad token exists (avoids warnings on generation)
+if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Wrap with a text-generation pipeline
+pipe = pipeline(
+    task="text-generation",
+    model=model,
+    tokenizer=tokenizer
+)
+
+def infer(prompt, max_new_tokens=128, temperature=0.7, top_p=0.9):
+    """Single-turn chat-style inference with Qwen 0.5B Instruct."""
+    if not prompt or not prompt.strip():
         return "Please type something to generate."
-    # Beam search is deterministic and tends to avoid loops on T5.
-    out = pipe(
-        prompt,
-        max_new_tokens=150,#int(max_new_tokens),
-        num_beams=5,#int(num_beams),
-        early_stopping=True,
-        length_penalty=1,#float(length_penalty),   # <1.0 favors shorter, reduces rambling
-        no_repeat_ngram_size=3                  # prevents 3-gram loops like "blue sky with"
+
+    # Use Qwen's chat template for better instruct-style behavior
+    messages = [
+        {"role": "system", "content": "You are a helpful, concise assistant for beginners learning about LLMs."},
+        {"role": "user", "content": prompt.strip()}
+    ]
+    chat_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True   # appends assistant prefix as the generation start
+    )
+
+    # Generation with light anti-repetition guards
+    outputs = pipe(
+        chat_prompt,
+        max_new_tokens=int(max_new_tokens),
+        do_sample=True,
+        temperature=float(temperature),
+        top_p=float(top_p),
+        no_repeat_ngram_size=3,      # prevents short n-gram loops
+        repetition_penalty=1.1,      # gentle nudge against repeating phrases
+        return_full_text=False       # only return the assistant's new text
     )
-    return out[0]["generated_text"]
+
+    return outputs[0]["generated_text"]
 
 demo = gr.Interface(
     fn=infer,
     inputs=[
-        gr.Textbox(lines=3, label="Instruction",
-                   placeholder="Explain in one paragraph: Why is the sky blue?"),
-        gr.Slider(32, 192, 96, step=8, label="Max new tokens"),
-        gr.Slider(2, 8, 4, step=1, label="Num beams"),
-        gr.Slider(0.6, 1.3, 0.9, step=0.05, label="Length penalty")
+        gr.Textbox(lines=3, label="Instruction", placeholder="Explain in one paragraph: Why is the sky blue?"),
+        gr.Slider(16, 256, 128, step=8, label="Max new tokens"),
+        gr.Slider(0.0, 1.5, 0.7, step=0.05, label="Temperature"),
+        gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-p"),
     ],
     outputs=gr.Textbox(lines=10, label="Output"),
-    title="FLAN-T5 (Deterministic) — Mini LLM"
+    title="Mini LLM (Local) — Qwen 2.5 (0.5B) Instruct"
 )
 
 if __name__ == "__main__":