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):
    return 'Привет'


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

demo.launch()