Create demo_t5_qa_pipe.py

demo_t5_qa_pipe.py

@@ -0,0 +1,39 @@
+import torch
+import tensorflow as tf
+import numpy as np
+from transformers import Text2TextGenerationPipeline
+
+class DemoT5QAPipeline(Text2TextGenerationPipeline):
+    def _forward(self, model_inputs, **generate_kwargs):
+        if self.framework == "pt":
+            in_b, input_length = model_inputs["input_ids"].shape
+        elif self.framework == "tf":
+            in_b, input_length = tf.shape(model_inputs["input_ids"]).numpy()
+
+        self.check_inputs(
+            input_length,
+            generate_kwargs.get("min_length", self.model.config.min_length),
+            generate_kwargs.get("max_length", self.model.config.max_length),
+        )
+        outputs = self.model.generate(**model_inputs, **generate_kwargs, return_dict_in_generate=True, output_scores=True, max_new_tokens=75)
+        
+        # Code from the parent class
+        output_ids = outputs.sequences
+        out_b = output_ids.shape[0]
+        if self.framework == "pt":
+            output_ids = output_ids.reshape(in_b, out_b // in_b, *output_ids.shape[1:])
+        elif self.framework == "tf":
+            output_ids = tf.reshape(output_ids, (in_b, out_b // in_b, *output_ids.shape[1:]))
+
+        output_sequences = outputs.sequences
+        output_scores = outputs.scores
+        return {"output_ids": output_ids, "output_sequences": output_sequences, "output_scores": output_scores}
+    
+    def postprocess(self, model_outputs):
+        guess_text = super().postprocess(model_outputs)[0]['generated_text']
+
+        transition_scores = self.model.compute_transition_scores(model_outputs['output_sequences'], model_outputs['output_scores'], normalize_logits=True)
+        log_probs = np.round(np.exp(transition_scores.cpu().numpy()), 3)[0]
+        guess_prob = np.product(log_probs)
+
+        return {'guess': guess_text, 'confidence': guess_prob}