from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer
from adapters import AutoAdapterModel
import os
import gradio as gr
import torch

HF_TOKEN = os.getenv("HF_TOKEN")
login(token=HF_TOKEN)

title = "Mental Health Chatbot"
description = "This bot is using a fine-tuned version of meta-llama/Llama-2-7b-chat-hf"

model_id = "meta-llama/Llama-2-7b-chat-hf"
adapter_model_id = "vojay/Llama-2-7b-chat-hf-mental-health"

model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)
model.load_adapter(adapter_model_id)


def predict(input, history=[]):
    new_user_input_ids = tokenizer.encode(f"{input}{tokenizer.eos_token}", return_tensors="pt")
    bot_input_ids = torch.cat([torch.LongTensor(history), new_user_input_ids], dim=-1)

    history = model.generate(
        bot_input_ids,
        max_length=4000,
        pad_token_id=tokenizer.eos_token_id
    ).tolist()

    response = tokenizer.decode(history[0]).split("<|endoftext|>")
    response = [
        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
    ]

    return response, history


gr.Interface(
    fn=predict,
    title=title,
    description=description,
    inputs=["text", "state"],
    outputs=["chatbot", "state"]
).launch()