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

model_id = "Atmanstr/qwen-financebot-finetuned"  # Replace with your HF repo path

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map=None,
    torch_dtype=torch.float16,
    trust_remote_code=True
)
model.eval()

def chat(message, history):
    system_prompt = "Your helpful assistant."
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Only return response part after last <|im_start|>assistant
    final_response = response.split("<|im_start|>assistant\n")[-1]
    return final_response

interface = gr.ChatInterface(fn=chat, title="Qwen 3B Chatbot", theme="soft")
interface.launch()