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

# --- Load model in 8-bit mode on CPU ---
print("Loading model in 8-bit on CPU (reduced memory)...")

quant_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_enable_fp32_cpu_offload=True,  # partial CPU offload
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "m-a-p/YuE-s1-7B-anneal-en-cot",
    trust_remote_code=True
)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    "m-a-p/YuE-s1-7B-anneal-en-cot",
    quantization_config=quant_config,
    device_map="cpu",        # force CPU usage
    torch_dtype=torch.float16,
    attn_implementation="eager",
    trust_remote_code=True
)

model.eval()
print("✅ Model loaded (8-bit, ~6–8 GB RAM).")

# --- Define generation function ---
def generate_text(prompt, max_tokens=200):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=int(max_tokens))
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# --- Gradio Interface ---
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Enter your prompt", lines=4, placeholder="Type something..."),
        gr.Slider(10, 400, value=200, step=10, label="Max tokens")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="YuE-s1-7B Text Generator (8-bit CPU)",
    description="A lightweight quantized version running on CPU (~6–8 GB RAM)."
)

if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7860)