import gradio as gr

from transformers import AutoModelForCausalLM, AutoTokenizer

@gr.cache(allow_output_mutation=True)
def load_model():
    model_id = "Tech-Meld/Hajax_Chat_1.0"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    return model, tokenizer

def get_response(input_text, model, tokenizer):
    inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
    outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

model, tokenizer = load_model()

def chat(input_text):
  response = get_response(input_text, model, tokenizer)
  return response

iface = gr.Interface(
    chat,
    "text",
    "text",
    title="Chat with Hajax",
    description="Type your message and press Enter to chat with the AI.",
)
iface.launch()