import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import os
HF_TOKEN = os.getenv('token')

MODEL_NAME = 'meta-llama/Llama-2-7b-chat-hf'
ADAPTERS_NAME = 'pivovalera2012/Llama-2-7b-Dr-Hous-test'

model_trained = AutoModelForCausalLM.from_pretrained(MODEL_NAME,
                                                    token=HF_TOKEN)

model_trained = PeftModel.from_pretrained(model_trained, ADAPTERS_NAME)
model_trained = model_trained.merge_and_unload()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)


def generate_text(prompt):
    encoding = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model_trained.generate(
            input_ids = encoding.input_ids,
            attention_mask = encoding.attention_mask,
            generation_config = generation_config
        )
    
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    answer = answer.split(':')
    return answer[1]


demo = gr.Interface(
    generate_text,
    inputs=["textbox"],
    outputs=["textbox"]
)

demo.launch()