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

base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
model = PeftModel.from_pretrained(base_model, "ewernn/perfect-refusal-model")
tokenizer = AutoTokenizer.from_pretrained("ewernn/perfect-refusal-model")

def chat(message, history):
    text = f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=64)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("model\n")[-1]
    return response.replace("<end_of_turn>", "").strip()

# Custom CSS to make interface more compact
css = """
.contain { max-height: 250px !important; }
.chatbot { min-height: 150px !important; max-height: 150px !important; }
footer { display: none !important; }
"""

demo = gr.ChatInterface(
    chat,
    css=css,
)

demo.launch()