from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from peft import PeftModel, PeftConfig
import torch
import gradio as gr

# Use the base model's ID
base_model_id = "mistralai/Mistral-7B-v0.1"
model_directory = "Tonic/mistralmed"

# Instantiate the Models
tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = 'left'

# Load the PEFT model
peft_config = PeftConfig.from_pretrained("Tonic/mistralmed")
base_model = AutoModelForSeq2SeqLM.from_pretrained(model_directory)
peft_model = PeftModel.from_pretrained(base_model, "Tonic/mistralmed")

class ChatBot:
    def __init__(self):
        self.history = []

    def predict(self, input):
        # Encode user input
        user_input_ids = tokenizer.encode(input + tokenizer.eos_token, return_tensors="pt")

        # Concatenate the user input with chat history
        if self.history:
            chat_history_ids = torch.cat([self.history, user_input_ids], dim=-1)
        else:
            chat_history_ids = user_input_ids

        # Generate a response using the PEFT model
        response = peft_model.generate(chat_history_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)

        # Update chat history
        self.history = response

        # Decode and return the response
        response_text = tokenizer.decode(response[0], skip_special_tokens=True)
        return response_text

bot = ChatBot()

title = "👋🏻Welcome to Tonic's MistralMed Chat🚀"
description = "You can use this Space to test out the current model (MistralMed) or duplicate this Space and use it for any other model on 🤗HuggingFace. Join me on Discord to build together."
examples = [["What is the boiling point of nitrogen"]]

iface = gr.Interface(
    fn=bot.predict,
    title=title,
    description=description,
    examples=examples,
    inputs="text",
    outputs="text",
    theme="ParityError/Anime"
)

iface.launch()