Update handler.py

handler.py

@@ -8,7 +8,7 @@ tqdm.pandas()
 from datasets import load_dataset
 from transformers import pipeline
 
-class EndpointHandler():
+class InferencePipeline():
     def __init__(self, path=""):
         df_wiki_windows = load_dataset("foxxy-hm/e2eqa-wiki",  data_files="processed/wikipedia_20220620_cleaned_v2.csv")["train"].to_pandas()
         df_wiki = load_dataset("foxxy-hm/e2eqa-wiki",  data_files="wikipedia_20220620_short.csv")["train"].to_pandas()
@@ -82,14 +82,11 @@ class EndpointHandler():
         return final_answer
 
     
-class InferencePipeline:
+class EndpointHandler():
     def __init__(self):
-        self.endpoint_handler = EndpointHandler()  # Instantiate the EndpointHandler class
-    
+        self.inference_pipeline = InferencePipeline() 
+        pipeline = pipeline("qa-model", model=self.inference_pipeline, tokenizer=None)
+        
     def __call__(self, question):
-        answer = self.endpoint_handler.get_answer_e2e(question)  # Call the get_answer_e2e method from EndpointHandler
-        return {"answer": answer}  # Return the answer as a dictionary
-
-
-inference_pipeline = InferencePipeline()
-pipeline = pipeline("qa-model", model=inference_pipeline, tokenizer=None)
+        answer = self.endpoint_handler.get_answer_e2e(question)  
+        return answer