Update src/deepeval/base_task.py

src/deepeval/base_task.py

@@ -1,274 +1,276 @@
-from abc import ABC, abstractmethod
-from datasets import load_dataset
-import os
-from dotenv import load_dotenv
-import openai
-from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList, LogitsProcessor
-import torch
-from typing import List
-from datetime import datetime
-load_dotenv()
-HF_TOKEN=os.getenv("HF_TOKEN")
-OPENAI_KEY = os.getenv("OPENAI_API_KEY")
-
-class BaseTask(ABC):
-    _model_cache = {}  # Class-level cache for models and tokenizers
-
-    def __init__(self, dataset_repo, model_name):
-        self.dataset_repo = dataset_repo
-        self.dataset = self.load_dataset_from_hf()
-
-        device_count = torch.cuda.device_count()
-        if device_count > 1:
-            self.device = "auto"
-            print(f"Using {device_count} GPUs with auto config.")
-        elif device_count == 1:
-            self.device = "cuda"
-            print(f"Using {device_count} GPU with cuda config.")
-        else:
-            self.device = "cpu"
-            print("No GPU found. Using CPU.")
-
-        self.model, self.tokenizer = self.get_cached_model(model_name, self.device)
-        openai.api_key = OPENAI_KEY
-
-
-    @classmethod
-    def get_cached_model(cls, model_name, device):
-        """Ensures the same model and tokenizer are used for every instance of subclasses."""
-        if model_name not in cls._model_cache:
-            cls._model_cache[model_name] = cls.load_model(model_name, device)
-        return cls._model_cache[model_name]
-    
-    @staticmethod
-    def load_model(model_name: str, device):
-        """Loads model and tokenizer once and caches it."""
-        print(f"Loading model: {model_name}")
-        start_time = datetime.now()
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map=device,
-            token=HF_TOKEN,  # Replace with actual token
-        )
-        end_time = datetime.now()
-        print(f"Model loaded in {(end_time - start_time).seconds} seconds.")
-        print("Model loaded.")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        return model, tokenizer
-
-    # @staticmethod
-    # def load_model(model_name: str, device, weight, dtype, base_model):
-    #     """Loads model and tokenizer once and caches it."""
-    #     print(f"Loading model: {model_name}")
-    #     start_time = datetime.now()
-    #     if weight == "Adapter":
-    #         base_model_1 = AutoModelForCausalLM.from_pretrained(
-    #             base_model,
-    #             torch_dtype=dtype,
-    #             device_map=device,
-    #             token=HF_TOKEN,  # Replace with actual token
-    #         )
-    #         model = PeftModel.from_pretrained(base_model_1, base_model)
-    #         tokenizer = AutoTokenizer.from_pretrained(base_model)
-    #         end_time = datetime.now()
-    #     else:
-    #         model = AutoModelForCausalLM.from_pretrained(
-    #             model_name,
-    #             torch_dtype=dtype,
-    #             device_map=device,
-    #             token=HF_TOKEN,  # Replace with actual token
-    #         )
-    #         tokenizer = AutoTokenizer.from_pretrained(model_name)
-    #         end_time = datetime.now()
-    #     print(f"Model loaded in {(end_time - start_time).seconds} seconds.")
-    #     print("Model loaded.")
-        
-    #     return model, tokenizer
-
-
-    def generate_response_mcqa(self, msg, max_new_tokens=1, choices: List[str]=[]):
-        # Ensure the tokenizer has a padding token
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token  # Use EOS token as PAD token
-
-        inputs = self.tokenizer(msg, return_tensors="pt", padding=True, truncation=True)
-        input_ids = inputs.input_ids
-        attention_mask = inputs.attention_mask
-
-        if self.model.config.pad_token_id is None:
-            self.model.config.pad_token_id = self.tokenizer.eos_token_id
-
-        # Get token IDs for answer choices
-        valid_answers = choices
-        valid_token_ids = [self.tokenizer.convert_tokens_to_ids(ans) for ans in valid_answers]
-
-        class MultipleChoiceLogitsProcessor:
-            def __call__(self, input_ids, scores):
-                mask = torch.full_like(scores, float("-inf"))
-                mask[:, valid_token_ids] = scores[:, valid_token_ids]  # Allow only valid tokens
-                return mask
-
-        logits_processor = LogitsProcessorList([MultipleChoiceLogitsProcessor()])
-
-        output = self.model.generate(
-            input_ids,
-            attention_mask=attention_mask,  # Fix: Pass attention_mask to avoid warning
-            max_new_tokens=max_new_tokens,
-            logits_processor=logits_processor
-        )
-        answer = self.tokenizer.decode(output[0][-1])
-
-        return answer
-    
-    def generate_response_mcqa_multi_token(self, msg, max_new_tokens=2, 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)
-
-        if self.device == "auto":
-            input_ids = inputs.input_ids
-            attention_mask = inputs.attention_mask
-        else:
-            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:
-
-        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.eos_token_id
-
-        chat = [
-            {"role": "user", "content": "You are a helpful AI assistant."},
-            {"role": "assistant", "content": "I am here to help you with any questions you may have."},
-            {"role": "user", "content": prompt},
-        ]
-        
-        formatted_chat = self.tokenizer.apply_chat_template(
-            chat,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-
-        inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
-
-        if self.device == "auto":
-            input_ids = inputs.input_ids
-            attention_mask = inputs.attention_mask
-        else:
-            input_ids = inputs.input_ids.to(self.model.device)
-            attention_mask = inputs.attention_mask.to(self.model.device)
-
-        output = self.model.generate(
-            input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=max_new_tokens,
-            do_sample=True,
-            temperature=0.7,
-        )
-
-        generated_ids = output[0]
-        prompt_len = input_ids.shape[1]
-        generated_tokens = generated_ids[prompt_len:]
-        result = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-        return result
-
-    def get_chat_template_tokens(self):
-        allowed_token_chat = [
-            {"role": "user", "content": ""},
-            {"role": "assistant", "content": ""}
-        ]
-        allowed_special_tokens = self.tokenizer.apply_chat_template(allowed_token_chat, tokenize=True)
-        return allowed_special_tokens
-
-    @abstractmethod
-    def load_dataset_from_hf(self):
-        """
-        Define your own loading method if needed.
-        :return: Dataset
-        """
-        print("Loading dataset from Hugging Face.")
-        start_time = datetime.now()
-        dataset= load_dataset(self.dataset_repo, token=HF_TOKEN, split="train")
-        print("Dataset loaded.")
-    
-        # Load 50 from each dataset
-        # if len(dataset) > 50:
-        #     dataset = dataset.shuffle(seed=42).select(range(50))
-        end_time = datetime.now()
-        print(f"Dataset loaded in {(end_time - start_time).seconds} seconds.")
-        return dataset
-    
-    def load_dataset_lmjudge_from_hf(self):
-        """
-        Define your own loading method if needed.
-        :return: Dataset
-        """
-        print("Loading dataset from Hugging Face.")
-        start_time = datetime.now()
-        dataset= load_dataset(self.dataset_repo, token=HF_TOKEN, split="train")
-        print("Dataset loaded.")
-        
-        #Load 100 from each dataset
-        if len(dataset) > 100:
-            dataset = dataset.shuffle(seed=42).select(range(100))
-        end_time = datetime.now()
-        print(f"Dataset loaded in {(end_time - start_time).seconds} seconds.")
-        return dataset
-
-    @abstractmethod
-    def evaluate(self):
+from abc import ABC, abstractmethod
+from datasets import load_dataset
+import os
+from dotenv import load_dotenv
+import openai
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessorList, LogitsProcessor
+import torch
+from typing import List
+from datetime import datetime
+load_dotenv()
+HF_TOKEN=os.getenv("HF_TOKEN")
+OPENAI_KEY = os.getenv("OPENAI_API_KEY")
+
+class BaseTask(ABC):
+    _model_cache = {}  # Class-level cache for models and tokenizers
+
+    def __init__(self, dataset_repo, model_name):
+        self.dataset_repo = dataset_repo
+        self.dataset = self.load_dataset_from_hf()
+
+        device_count = torch.cuda.device_count()
+        if device_count > 1:
+            self.device = "auto"
+            print(f"Using {device_count} GPUs with auto config.")
+        elif device_count == 1:
+            self.device = "cuda"
+            print(f"Using {device_count} GPU with cuda config.")
+        else:
+            self.device = "cpu"
+            print("No GPU found. Using CPU.")
+
+        self.model, self.tokenizer = self.get_cached_model(model_name, self.device)
+        openai.api_key = OPENAI_KEY
+
+
+    @classmethod
+    def get_cached_model(cls, model_name, device):
+        """Ensures the same model and tokenizer are used for every instance of subclasses."""
+        if model_name not in cls._model_cache:
+            cls._model_cache[model_name] = cls.load_model(model_name, device)
+        return cls._model_cache[model_name]
+    
+    @staticmethod
+    def load_model(model_name: str, device):
+        """Loads model and tokenizer once and caches it."""
+        print(f"Loading model: {model_name}")
+        start_time = datetime.now()
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map=device,
+            token=HF_TOKEN,  # Replace with actual token
+        )
+        end_time = datetime.now()
+        print(f"Model loaded in {(end_time - start_time).seconds} seconds.")
+        print("Model loaded.")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        return model, tokenizer
+
+    # @staticmethod
+    # def load_model(model_name: str, device, weight, dtype, base_model):
+    #     """Loads model and tokenizer once and caches it."""
+    #     print(f"Loading model: {model_name}")
+    #     start_time = datetime.now()
+    #     if weight == "Adapter":
+    #         base_model_1 = AutoModelForCausalLM.from_pretrained(
+    #             base_model,
+    #             torch_dtype=dtype,
+    #             device_map=device,
+    #             token=HF_TOKEN,  # Replace with actual token
+    #         )
+    #         model = PeftModel.from_pretrained(base_model_1, base_model)
+    #         tokenizer = AutoTokenizer.from_pretrained(base_model)
+    #         end_time = datetime.now()
+    #     else:
+    #         model = AutoModelForCausalLM.from_pretrained(
+    #             model_name,
+    #             torch_dtype=dtype,
+    #             device_map=device,
+    #             token=HF_TOKEN,  # Replace with actual token
+    #         )
+    #         tokenizer = AutoTokenizer.from_pretrained(model_name)
+    #         end_time = datetime.now()
+    #     print(f"Model loaded in {(end_time - start_time).seconds} seconds.")
+    #     print("Model loaded.")
+        
+    #     return model, tokenizer
+
+
+    def generate_response_mcqa(self, msg, max_new_tokens=1, choices: List[str]=[]):
+        # Ensure the tokenizer has a padding token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token  # Use EOS token as PAD token
+
+        inputs = self.tokenizer(msg, return_tensors="pt", padding=True, truncation=True)
+        input_ids = inputs.input_ids
+        attention_mask = inputs.attention_mask
+
+        if self.model.config.pad_token_id is None:
+            self.model.config.pad_token_id = self.tokenizer.eos_token_id
+
+        # Get token IDs for answer choices
+        valid_answers = choices
+        valid_token_ids = [self.tokenizer.convert_tokens_to_ids(ans) for ans in valid_answers]
+
+        class MultipleChoiceLogitsProcessor:
+            def __call__(self, input_ids, scores):
+                mask = torch.full_like(scores, float("-inf"))
+                mask[:, valid_token_ids] = scores[:, valid_token_ids]  # Allow only valid tokens
+                return mask
+
+        logits_processor = LogitsProcessorList([MultipleChoiceLogitsProcessor()])
+
+        output = self.model.generate(
+            input_ids,
+            attention_mask=attention_mask,  # Fix: Pass attention_mask to avoid warning
+            max_new_tokens=max_new_tokens,
+            logits_processor=logits_processor
+        )
+        answer = self.tokenizer.decode(output[0][-1])
+
+        return answer
+    
+    def generate_response_mcqa_multi_token(self, msg, max_new_tokens=2, 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)
+
+        if self.device == "auto":
+            input_ids = inputs.input_ids
+            attention_mask = inputs.attention_mask
+        else:
+            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:
+
+        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.eos_token_id
+
+        chat = [
+            {"role": "user", "content": "You are a helpful AI assistant."},
+            {"role": "assistant", "content": "I am here to help you with any questions you may have."},
+            {"role": "user", "content": prompt},
+        ]
+        
+        formatted_chat = self.tokenizer.apply_chat_template(
+            chat,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
+
+        if self.device == "auto":
+            input_ids = inputs.input_ids
+            attention_mask = inputs.attention_mask
+        else:
+            input_ids = inputs.input_ids.to(self.model.device)
+            attention_mask = inputs.attention_mask.to(self.model.device)
+
+        output = self.model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+        )
+
+        generated_ids = output[0]
+        prompt_len = input_ids.shape[1]
+        generated_tokens = generated_ids[prompt_len:]
+        result = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+        return result
+
+    def get_chat_template_tokens(self):
+        allowed_token_chat = [
+            {"role": "user", "content": ""},
+            {"role": "assistant", "content": ""}
+        ]
+        allowed_special_tokens = self.tokenizer.apply_chat_template(allowed_token_chat, tokenize=True)
+        return allowed_special_tokens
+
+    @abstractmethod
+    def load_dataset_from_hf(self):
+        """
+        Define your own loading method if needed.
+        :return: Dataset
+        """
+        print("Loading dataset from Hugging Face.")
+        start_time = datetime.now()
+        dataset= load_dataset(self.dataset_repo, token=HF_TOKEN, split="train")
+        print("Dataset loaded.")
+    
+        # Load 50 from each dataset
+        mcqa_sample_size = 3
+        if len(dataset) > mcqa_sample_size:
+            dataset = dataset.shuffle(seed=42).select(range(mcqa_sample_size))
+        end_time = datetime.now()
+        print(f"Dataset loaded in {(end_time - start_time).seconds} seconds.")
+        return dataset
+    
+    def load_dataset_lmjudge_from_hf(self):
+        """
+        Define your own loading method if needed.
+        :return: Dataset
+        """
+        print("Loading dataset from Hugging Face.")
+        start_time = datetime.now()
+        dataset= load_dataset(self.dataset_repo, token=HF_TOKEN, split="train")
+        print("Dataset loaded.")
+        
+        #Load 100 from each dataset
+        llmjudge_sample_size = 3
+        if len(dataset) > llmjudge_sample_size:
+            dataset = dataset.shuffle(seed=42).select(range(llmjudge_sample_size))
+        end_time = datetime.now()
+        print(f"Dataset loaded in {(end_time - start_time).seconds} seconds.")
+        return dataset
+
+    @abstractmethod
+    def evaluate(self):
         pass