Upload TestEnsembleQAPipeline

config.json

@@ -6,7 +6,7 @@
   "classifier_dropout": 0.0,
   "custom_pipelines": {
     "test-qa": {
-      "impl": "__main__.TestEnsembleQAPipeline",
+      "impl": "test_ensemble.TestEnsembleQAPipeline",
       "pt": [
         "AutoModelForSeq2SeqLM"
       ],

test_ensemble.py

@@ -9,25 +9,49 @@ from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import gzip
 from transformers import Text2TextGenerationPipeline
+from transformers import Text2TextGenerationPipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+import numpy as np
 
 class TestEnsembleQAPipeline(Text2TextGenerationPipeline):
-    def __init__(self, model=None, tokenizer=None, framework="pt", **kwargs):  # Add **kwargs to accept additional unused arguments
+    def __init__(self, model=None, tokenizer=None, framework="pt", **kwargs):
         super().__init__(model=model, tokenizer=tokenizer, framework=framework)
-        self.quiz_bowl_model = QuizBowlModel()  # Initializes your QuizBowl model
+        self.quiz_bowl_model = QuizBowlModel()  # Initialize your QuizBowl model
+
+    def preprocess(self, text, **kwargs):
+        """Prepare the text inputs for processing."""
+        prompt = "Please provide a concise answer to the following question:"
+        input_text = f"{prompt} {text}"
+        return self.tokenizer(input_text, return_tensors=self.framework, padding=True, truncation=True)
 
     def _forward(self, model_inputs, **generate_kwargs):
-        questions = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in model_inputs["input_ids"]]
-        results = self.quiz_bowl_model.guess_and_buzz(questions)
-        return results
+        """Forward pass to generate outputs from the model."""
+        if self.framework == "pt":
+            model_outputs = self.model.generate(**model_inputs, **generate_kwargs, return_dict_in_generate=True, output_scores=True)
+        else:
+            raise NotImplementedError("TensorFlow framework is not supported in this pipeline.")
+        return model_outputs
 
     def postprocess(self, model_outputs):
-        results = {}
-        for output in model_outputs:
-            guess_text = output['guess']
-            confidence = output['confidence']
-            results = {'guess': guess_text, 'confidence': confidence}
+        """Process model outputs to extract answers and confidence scores."""
+        results = []
+        for output in model_outputs.sequences:
+            decoded_text = self.tokenizer.decode(output, skip_special_tokens=True)
+            scores = self.calculate_confidence(model_outputs.scores)
+            results.append({'guess': decoded_text, 'confidence': scores})
         return results
 
+    def calculate_confidence(self, scores):
+        """Calculate confidence from the model's score outputs."""
+        if scores:
+            log_probs = [torch.nn.functional.log_softmax(score, dim=-1) for score in scores]
+            avg_scores = [log_probs[i][0, output[i + 1]].item() for i in range(len(output) - 1)]
+            confidence_score = np.exp(np.mean(avg_scores))
+        else:
+            confidence_score = None
+        return confidence_score
+
+
 class QuizBowlModel:
     def __init__(self):
         self.load_models()
@@ -55,10 +79,10 @@ class QuizBowlModel:
         """Generate answers from all models for given questions"""
         total_answers = self.generate_answers(question_texts)
         # Display the model's guesses before voting
-        print("Answers Before Voting Mechanism:")
+        # print("Answers Before Voting Mechanism:")
 
-        for question, model_answers in zip(question_texts, total_answers):
-            print(f"{question}\nModel Guesses: {model_answers}\n\n")
+        # for question, model_answers in zip(question_texts, total_answers):
+        #     print(f"{question}\nModel Guesses: {model_answers}\n\n")
         return self.ensemble_tfidf_voting(total_answers)
 
     def generate_answers(self, question_texts):