app.py

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

model_id = "PowerInfer/SmallThinker-21BA3B-Instruct"
model = None
tokenizer = None
generator = None

def chat(prompt, max_new_tokens=256, temperature=0.7):
    global model, tokenizer, generator
    if generator is None:
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",  # Let Accelerate handle the devices
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            trust_remote_code=True
        )
        # ❌ Remove the `device` argument to avoid ValueError
        generator = pipeline(
            "text-generation",
            model=model,
            tokenizer=tokenizer
        )

    output = generator(
        prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )
    return output[0]["generated_text"]

demo = gr.Interface(
    fn=chat,
    inputs=[
        gr.Textbox(label="Prompt", lines=4),
        gr.Slider(32, 512, value=256, step=16, label="Max New Tokens"),
        gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
    ],
    outputs=gr.Textbox(label="Response"),
    title="💬 SmallThinker-21B",
)

if __name__ == "__main__":
    demo.launch()