import gradio as gr
import joblib
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer

# Load the model
loaded_model = joblib.load("qa_model.joblib")
vectorizer = loaded_model["vectorizer"]
tfidf_matrix = loaded_model["tfidf_matrix"]
paragraphs = loaded_model["paragraphs"]

def answer_question(question):
    question_vector = vectorizer.transform([question])
    similarities = cosine_similarity(question_vector, tfidf_matrix)
    most_similar_paragraph_index = np.argmax(similarities)
    most_similar_paragraph = paragraphs[most_similar_paragraph_index]

    paragraph_sentences = most_similar_paragraph.split(".")
    best_sentence = ""
    max_overlap = 0

    question_words = set(question.lower().split())

    for sentence in paragraph_sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
        sentence_words = set(sentence.lower().split())
        overlap = len(question_words.intersection(sentence_words))
        if overlap > max_overlap:
            max_overlap = overlap
            best_sentence = sentence

    return best_sentence.strip()

iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question here..."),
    outputs="text",
    title="Mahabharata Question Answering",
    description="Ask a question about the Mahabharata, and the model will attempt to answer it.",
)

iface.launch()