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

MODEL = "mistralai/Mistral-7B-Instruct-v0.3"

tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16,
    device_map="auto"
)

def chat(message, history):
    messages = [
        {"role": "user", "content": message}
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        return_tensors="pt",
        add_generation_prompt=True
    ).to(model.device)

    outputs = model.generate(
        inputs,
        max_new_tokens=256
    )

    reply = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    return reply

demo = gr.ChatInterface(chat)

demo.launch(
    server_name="0.0.0.0",
    server_port=7860
)