Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForQuestionAnswering
 import torch
-from mlcroissant import Dataset
+from datasets import load_dataset
 import random
 
 # Load the DistilBERT model and tokenizer
@@ -10,27 +10,27 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 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"))
+ds = load_dataset("knowrohit07/gita_dataset")
 
 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])
+    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")
-    input_ids = inputs["input_ids"].tolist()[0]
-
+    inputs = tokenizer.encode_plus(question, context, add_special_tokens=True, return_tensors="pt", max_length=512, truncation=True)
+    
     # Get the answer
-    outputs = model(**inputs)
+    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(input_ids[answer_start:answer_end]))
+    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]":
@@ -41,13 +41,9 @@ def generate_response(question):
     
     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=predict,
+    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",
@@ -61,5 +57,5 @@ iface = gr.Interface(
     ]
 )
 
-# Launch the interface with sharing enabled
-iface.launch(share=True)
+# Launch the interface
+iface.launch()