Re added function that got lost with commit

src/deepeval/base_task.py

@@ -71,6 +71,62 @@ class BaseTask(ABC):
         answer = self.tokenizer.decode(output[0][-1])
 
         return answer
+    
+    def generate_response_mcqa_multi_token(self, msg, max_new_tokens=5, choices: list = []):
+        """
+        Handles multiple-choice questions where answers might have multiple tokens.
+        """
+        # Ensure tokenizer has proper special tokens set
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        if self.model.config.pad_token_id is None:
+            self.model.config.pad_token_id = self.tokenizer.pad_token_id
+
+        chat = [
+                {"role": "user", "content": "You are a multiple choice question-answering chatbot. Do not give an answer that is not included in the choices. Only answer with letters like A, B, C, D..."},
+                {"role": "assistant", "content": "I am ready to answer your questions. Feel free to ask anything.\n"},
+                {"role": "user", "content": f"{msg}"},
+            ]
+        formatted_chat = self.tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+        print(formatted_chat)
+        inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
+        input_ids = inputs.input_ids.to(self.model.device)
+        attention_mask = inputs.attention_mask.to(self.model.device)
+
+        # Generate the sequence of letters starting from 'A'
+        letters = [chr(ord('A') + i) for i in range(len(choices))]  # Create option letters A, B, C, D, E, ...
+        encoded_choices = [self.tokenizer.encode(letter, add_special_tokens=False) for letter in letters]
+        flattened_encoded_choices = [item for sublist in encoded_choices for item in sublist]  # Flatten the list
+        print(flattened_encoded_choices)
+
+        allowed_tokens = flattened_encoded_choices
+        allowed_tokens += self.get_chat_template_tokens() # Get the special chat tokens
+        allowed_token_ids = set(allowed_tokens)  # Ensure uniqueness
+
+        # Custom LogitsProcessor to restrict generation
+        class RestrictToABCDLogitsProcessor(LogitsProcessor):
+            def __call__(self, input_ids, scores):
+                mask = torch.full_like(scores, float("-inf"))  # Block all tokens
+                mask[:, list(allowed_token_ids)] = scores[:, list(allowed_token_ids)]  # Allow only A, B, C, D tokens
+                return mask
+        logits_processor = LogitsProcessorList([RestrictToABCDLogitsProcessor()])
+
+        # Generate response
+        output = self.model.generate(
+            input_ids,
+            do_sample=True,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
+            temperature=0.4,
+            logits_processor=logits_processor,
+        )
+        generated_ids = output[0]  # The generated sequence including the prompt
+        generated_tokens = generated_ids[len(input_ids[0]):]  # Exclude the input_ids part
+        generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        return generated_text
 
     def generate_response(self, prompt: str, max_new_tokens: int = 100) -> str: