Update app.py

app.py

@@ -1,33 +1,65 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import torch
+from mlcroissant import Dataset
+import random
 
+# Load the DistilBERT model and tokenizer
 model_name = "distilbert/distilbert-base-cased-distilled-squad"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-gita_context = """
-The Bhagavad Gita is a 700-verse Hindu scripture that is part of the Indian epic Mahabharata. It is a dialogue between Prince Arjuna and Lord Krishna, who serves as his charioteer. The Gita's core message includes:
-1. The immortality of the soul (Atman)
-2. The nature of action (Karma) and duty (Dharma)
-3. The importance of devotion (Bhakti)
-4. The pursuit of knowledge (Jnana) and wisdom
-5. Different types of Yoga: Karma Yoga, Bhakti Yoga, Jnana Yoga, and Raja Yoga
-6. The concept of detachment from the fruits of one's actions
-7. The divine nature of Krishna as an avatar of Vishnu
-Key teachings include performing one's duty without attachment to results, the importance of self-realization, and the path to liberation (Moksha).
-"""
+model = AutoModelForQuestionAnswering.from_pretrained(model_name)
+
+# Load the Bhagavad Gita dataset
+ds = Dataset(jsonld="https://huggingface.co/api/datasets/knowrohit07/gita_dataset/croissant")
+records = list(ds.records("default"))
+
+def get_relevant_context(question):
+    # Randomly select 5 records to form the context
+    selected_records = random.sample(records, 5)
+    context = " ".join([record["Text"] for record in selected_records])
+    return context
 
 def generate_response(question):
-    prompt = f"Based on the following context about the Bhagavad Gita, answer the question.\n\nContext: {gita_context}\n\nQuestion: {question}\n\nAnswer:"
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    outputs = model.generate(input_ids, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.95)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    context = get_relevant_context(question)
+    
+    # Encode the question and context
+    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt")
+    input_ids = inputs["input_ids"].tolist()[0]
+
+    # Get the answer
+    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(input_ids[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
+
+# Define the predict function for the API
+def predict(question):
+    return generate_response(question)
 
+# Create the Gradio interface
 iface = gr.Interface(
-    fn=generate_response,
+    fn=predict,
     inputs=gr.Textbox(lines=2, placeholder="Enter your question about the Bhagavad Gita here..."),
-    outputs="text"
+    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?"]
+    ]
 )
 
-iface.launch()
+# Launch the interface with sharing enabled
+iface.launch(share=True)