Upload MyTestPipeline

new_task.py

@@ -2,10 +2,11 @@ from transformers import Text2TextGenerationPipeline, AutoModelForSeq2SeqLM, TFA
 import torch
 import tensorflow as tf
 import numpy as np
+import math
 
 class MyTestPipeline(TextGenerationPipeline):
     def preprocess(self, text, **kwargs):
-        prompt = 'Answer the following question/statement without any explanation, do not abbreviate names.'
+        prompt = 'Answer the following question/statement in English without any explanation, do not abbreviate names.'
         txt = f"<|user|>\n{prompt} {text}\n<|end|>\n<|assistant|>"
         return self.tokenizer(txt, return_tensors=self.framework)
     
@@ -15,7 +16,7 @@ class MyTestPipeline(TextGenerationPipeline):
         elif self.framework == "tf":
             in_b, input_length = tf.shape(model_inputs["input_ids"]).numpy()
 
-        outputs = self.model.generate(**model_inputs, **generate_kwargs, return_dict_in_generate=True, output_scores=True, max_new_tokens=75)
+        outputs = self.model.generate(**model_inputs, **generate_kwargs, return_dict_in_generate=True, output_scores=True)
 
         output_ids = outputs.sequences
         out_b = output_ids.shape[0]
@@ -35,8 +36,7 @@ class MyTestPipeline(TextGenerationPipeline):
         log_probs = np.round(np.exp(transition_scores.cpu().numpy()), 3)[0]
         guess_prob = np.product(log_probs)
 
-        if guess_prob >= 0.8:
-          guess_prob = 1.0
+        confidence = (math.exp(12*(guess_prob - 0.5))) / (1 + math.exp(12 * (guess_prob - 0.5)))
 
-        return {'guess': guess_text, 'confidence': guess_prob}
+        return {'guess': guess_text, 'confidence': confidence}