Added time stamps for cost analysis

src/deepeval/base_task.py

@@ -6,6 +6,7 @@ import openai
 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")
@@ -32,12 +33,15 @@ class BaseTask(ABC):
     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.")
         print("Model loaded.")
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         return model, tokenizer
@@ -186,6 +190,7 @@ class BaseTask(ABC):
         :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.")
         
@@ -193,6 +198,8 @@ class BaseTask(ABC):
         print("Original dataset size: ", len(dataset))
         dataset = dataset.shuffle(seed=42).select(range(int(len(dataset) * 0.25)))
         print("Reduced dataset size: ", len(dataset))
+        end_time = datetime.now()
+        print(f"Dataset loaded in {end_time - start_time} seconds.")
         return dataset
 
     @abstractmethod

src/deepeval/deepeval_task_manager.py

@@ -15,6 +15,7 @@ from src.deepeval.complex_reasoning import ComplexReasoningTask
 from src.deepeval.truthfulness_task import TruthfulnessTask
 from src.deepeval.nli import NLITask
 from typing import List
+from datetime import datetime
 load_dotenv()
 HF_TOKEN=os.getenv("HF_TOKEN")
 
@@ -59,15 +60,21 @@ class DeepEvalTaskManager:
     def run_tasks(self):
         """Execute validated tasks in order."""
         results = {}
+        total_start_time = datetime.now()
         for task_name, task_method in self.tasks_to_run.items():
             try:
+                start_time = datetime.now()
                 print("Running task: ", task_name)
                 task_enum = getattr(Task, task_name)
                 task_value = task_enum.value
                 results[task_value] = task_method()  # Call the stored method reference
+                end_time = datetime.now()
+                print(f"Task {task_name} completed in {end_time - start_time} seconds.")
             except Exception as e:
                 print(f"Error At Task: {task_name} - {e}")
                 continue
+        total_end_time = datetime.now()
+        print(f"All tasks completed in {total_end_time - total_start_time} seconds.")
         print("All tasks completed.")
         return results