Update app.py

app.py

@@ -14,6 +14,7 @@ import gradio as gr
 import pandas as pd
 from spacy.cli import download
 
+
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
@@ -23,10 +24,39 @@ class BasicAgent:
             download("en_core_web_sm")
         self.spacy = spacy.load("en_core_web_sm")
         self.whisper_model = whisper.load_model("base")
-        self.qa_pipeline = pipeline("question-answering")
+        self.qa_pipeline = pipeline("question-answering", truncation=True, padding=True)
         self.ner_pipeline = pipeline("ner", aggregation_strategy="simple")
-        self.embedding_model = pipeline("feature-extraction")
+        self.embedding_model = pipeline("feature-extraction", truncation=True)
+
+    def split_text_into_chunks(self, text, max_length=512):
+        """Split text into chunks smaller than `max_length` tokens."""
+        words = text.split()
+        chunks = []
+        chunk = []
+
+        for word in words:
+            chunk.append(word)
+            if len(' '.join(chunk)) > max_length:
+                chunks.append(' '.join(chunk[:-1]))  # Add the chunk and reset
+                chunk = [word]
+
+        if chunk:
+            chunks.append(' '.join(chunk))  # Add the final chunk
+
+        return chunks
 
+    def answer_question(self, question: str, context: str) -> str:
+        try:
+            context_chunks = self.split_text_into_chunks(context, max_length=512)
+            answers = []
+            for chunk in context_chunks:
+                answer = self.qa_pipeline(question=question, context=chunk)["answer"]
+                answers.append(answer)
+
+            return " ".join(answers)  # Combine answers from chunks
+        except Exception as e:
+            return f"Error answering question: {e}"
+            
     def extract_named_entities(self, text):
         entities = self.ner_pipeline(text)
         return [e["word"] for e in entities if e["entity_group"] == "PER"]