import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig

# Your new repository!
peft_model_id = "tenith/my_qwen"

print("Loading config...")
config = PeftConfig.from_pretrained(peft_model_id)

print("Loading base model onto CPU...")
base_model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    trust_remote_code=True,
    device_map="cpu",
    torch_dtype=torch.float32 
)

print("Loading custom adapters...")
model = PeftModel.from_pretrained(base_model, peft_model_id)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id, trust_remote_code=True)

def generate_response(message, history):
    inputs = tokenizer(message, return_tensors="pt").to("cpu")
    
    outputs = model.generate(
        input_ids=inputs["input_ids"], 
        attention_mask=inputs["attention_mask"],
        max_new_tokens=150,
        pad_token_id=tokenizer.eos_token_id
    )
    
    raw_answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    clean_answer = raw_answer.replace(message, "").strip()
    return clean_answer

demo = gr.ChatInterface(
    fn=generate_response,
    title="Bitcoin Computer Assistant",
    description="Ask me anything about Bitcoin Computer!",
    examples=["What is Bitcoin Computer?", "How do you create a Computer instance?"]
)

demo.launch()