Update hybrid_pipe.py

hybrid_pipe.py

@@ -61,8 +61,9 @@ class HybridQAModel(nn.Module, PyTorchModelHubMixin):
         
     def infer_generative(self, model, tokenizer, input_text, **generate_kwargs):
         max_input_length = min(tokenizer.model_max_length, model.config.max_length)
+        input_text += " Do not output anything but the question's answer."
         messages = [
-          {"role": "user", "content": input_text + " Do not output anything but the question's answer."}
+          {"role": "user", "content": input_text}
         ]
         input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
         generated_outputs = model.generate(input_ids, max_new_tokens=256, temperature=0.5, output_scores=True, return_dict_in_generate=True)
@@ -74,7 +75,7 @@ class HybridQAModel(nn.Module, PyTorchModelHubMixin):
         average_confidence = sum(max_confidence_scores) / len(max_confidence_scores)  # Calculate average confidence
 
         decoded_output = tokenizer.decode(generated_outputs.sequences[0], skip_special_tokens=True)
-        final_output = decoded_output.replace("<|im_end|>", "").split("\n")[-1]
+        final_output = decoded_output[len(input):].split("\n")[-1]
         average_confidence, final_output
         return final_output, average_confidence