import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_ID = "Qwen/Qwen2.5-Coder-3B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="cpu"
)

# IMPORTANT: set pad token
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

def generate(prompt: str):
    if not prompt.strip():
        return "⚠️ Empty prompt"

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        padding=True
    )

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_new_tokens=128,
            do_sample=False,
            use_cache=False,          # 🔥 THIS AVOIDS THE BUG
            pad_token_id=tokenizer.eos_token_id
        )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=8, placeholder="Enter prompt"),
    outputs=gr.Textbox(lines=12),
    title="Qwen 2.5 Coder – Stable CPU Mode"
).launch()