import gradio as gr
from datasets import load_dataset
from transformers import pipeline

# Загружаем твой dataset
dataset = load_dataset("adiom/ai-dataset")
data = dataset['train']

# Загружаем модель для вопросов и ответов
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")

# Функция для обработки запросов
def answer_question(query, confidence_threshold=0.1):
    best_answer = None
    highest_score = 0
    
    # Проходим по всем текстам в dataset и находим лучший ответ
    for item in data:
        context = item['text']  # Используем 'text' из dataset
        result = qa_pipeline(question=query, context=context)
        
        # Если уверенность выше порога, обновляем лучший ответ
        if result["score"] > confidence_threshold and result["score"] > highest_score:
            best_answer = result["answer"]
            highest_score = result["score"]

    # Если ответа нет, возвращаем сообщение
    if best_answer is None:
        return "Извините, я не нашел подходящего ответа."
    
    return best_answer

# Интерфейс Gradio
iface = gr.Interface(fn=answer_question, 
                     inputs=["text", gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.1, label="Confidence Threshold")], 
                     outputs="text",
                     live=True,
                     title="Canfly Tiny-Inna",
                     description="Задайте вопрос о Canfly")

iface.launch(share=True)