import gradio as gr
import os
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
access_token=os.environ['token']

# Load Model
model_directory = 'paulpall/GEC_Estonian_OPUS-MT'
tokenizer = AutoTokenizer.from_pretrained(model_directory, token=access_token)
model = AutoModelForSeq2SeqLM.from_pretrained(model_directory, token=access_token)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    # Generate corrected sentence
    input_ids = tokenizer.encode(message, padding='max_length', truncation=True, max_length=128, return_tensors='pt')
    output_ids = model.generate(input_ids=input_ids.to(model.device))
    output_sentence = tokenizer.decode(output_ids[0], skip_special_tokens=True).replace(r"▁",r" ")

    response = output_sentence

    yield response

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond
)

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