import gradio as gr
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import torch
from datasets import load_dataset
import random

# Load the DistilBERT model and tokenizer
model_name = "distilbert/distilbert-base-cased-distilled-squad"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# Load the Bhagavad Gita dataset
ds = load_dataset("knowrohit07/gita_dataset")

def get_relevant_context(question):
    # Randomly select 5 records to form the context
    selected_records = random.sample(ds['train'], 5)
    context = " ".join([record['Text'] for record in selected_records])
    return context

def generate_response(question):
    context = get_relevant_context(question)
    
    # Encode the question and context
    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt", max_length=512, truncation=True)
    
    # Get the answer
    with torch.no_grad():
        outputs = model(**inputs)
    
    answer_start = torch.argmax(outputs.start_logits)
    answer_end = torch.argmax(outputs.end_logits) + 1
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0][answer_start:answer_end]))

    # If the model couldn't find an answer, provide a default response
    if answer == "" or answer == "[CLS]" or answer == "[SEP]":
        answer = "I'm sorry, but I couldn't find a specific answer to that question in the Bhagavad Gita. Could you please rephrase your question or ask about a different topic from the Gita?"

    # Add a disclaimer
    disclaimer = "\n\nPlease note: This response is generated by an AI model based on the Bhagavad Gita dataset. For authoritative information, please consult the original text or scholarly sources."
    
    return answer + disclaimer

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your question about the Bhagavad Gita here..."),
    outputs="text",
    title="Bhagavad Gita Q&A Assistant",
    description="Ask questions about the Bhagavad Gita. The AI will attempt to provide answers based on the text.",
    examples=[
        ["What is the main message of the Bhagavad Gita?"],
        ["Who is Krishna in the Bhagavad Gita?"],
        ["What does the Gita say about dharma?"],
        ["How does the Bhagavad Gita define yoga?"],
        ["What is the significance of Arjuna's dilemma?"]
    ]
)

# Launch the interface
iface.launch()