import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# load the Xortron model
MODEL_ID = "darkc0de/XortronCriminalComputingConfig"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

def respond(message, history):
    inputs = tokenizer(message, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )
    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return reply

demo = gr.ChatInterface(
    fn=respond,
    type="messages",  # avoids that deprecation warning
    chatbot=gr.Chatbot(height=600, show_copy_button=True),
    textbox=gr.Textbox(placeholder="Chat with Xortron...", container=False, scale=7),
    title="Xortron Chat",
)

if __name__ == "__main__":
    demo.launch()