customize the code for peft model

src/deepeval/base_task.py

@@ -3,6 +3,7 @@ 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
@@ -30,20 +31,32 @@ class BaseTask(ABC):
         return cls._model_cache[model_name]
 
     @staticmethod
-    def load_model(model_name: str, device):
+    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()
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map=device,
-            token=HF_TOKEN,  # Replace with actual token
-        )
-        end_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.")
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
         return model, tokenizer