import gradio as gr
from src.retriever import CustomRetriever
from src.generator import CustomGenerator
from src.model import MyRagModel

retriever = CustomRetriever(r'data/chunks.json', r'data/embeddings.index', r'data/metadata.json')
generator = CustomGenerator()
rag = MyRagModel(retriever, generator)


def generate_response(query):
    return rag.get_answer(query)


interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label='Введите ваш вопрос: '),
    outputs=gr.Textbox(label='Ответ модели: '),
    title='Rag Model Demo',
    description='QA ассистент по вопросам законов Кыргызской Республики. Может отвечать на вопросы по типу: '
                '"Сколько лет дают за убийство человека?"',
    allow_flagging='never',
)


if __name__ == '__main__':
    interface.launch(share=True)