import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model and tokenizer
model_name = "MajorJalud/Qwen2.5-0.5B-Instruct-Gensyn-Swarm-reptilian_strong_gull"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)


# Function to handle the chat
def chat(message, history):
    # Add a prompt to make the model act like a helpful assistant
    inputs = tokenizer("You are a helpful assistant. User: " + message, return_tensors="pt")
    prompt = f"You are a helpful assistant. User: {message} Assistant: "
    inputs = tokenizer(prompt, return_tensors="pt")
    # Generate a response
    outputs = model.generate(**inputs, max_length=100, pad_token_id=tokenizer.eos_token_id)
    # Decode the response to text
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Remove the prompt part from the response
    response = response.replace(prompt, "")
    # Return the conversation
    return [(message, response)]

# Create a chat interface
gr.ChatInterface(chat).launch()