src/evaluators/evaluate.py

import json
import os
from datetime import datetime,timedelta,timezone
from typing import Dict
from dataclasses import dataclass
from enum import Enum
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer,AutoModel
import traceback

from src.evaluators import EVALUATOR_REGISTRY
from src.evaluators.base_evaluator import BaseEvaluator
from src.envs import API, EVAL_REQUESTS_PATH, RESULTS_REPO, QUEUE_REPO,TOKEN


class EvaluationStatus(Enum):
    PENDING = "PENDING"
    RUNNING = "RUNNING"
    FINISHED = "FINISHED"
    FAILED = "FAILED"

@dataclass
class EvaluationResult:
    """Dataclass to hold the results of a single model evaluation."""
    model: str
    revision: str
    precision: str
    weight_type: str
    results: Dict[str, float]
    error: str = None




def evaluate_model(model_name: str, revision: str, precision: str, weight_type: str) -> EvaluationResult:
    """
    Evaluates a model on ALL registered tasks.
    """
    try:
        print(f"\nStarting evaluation for model: {model_name}")
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load model & tokenizer ONCE
        print("Loading classification model and tokenizer...")
        classification_model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            revision=revision,
            torch_dtype=getattr(torch, precision),
            trust_remote_code=True
        ).to(device)
        tokenizer = AutoTokenizer.from_pretrained(model_name, revision=revision)
        print("✅ Classification Model loaded successfully.")
        print("Loading base model...")
        embdding_model = AutoModel.from_pretrained(
            model_name,
            revision=revision,
            torch_dtype=getattr(torch, precision),
            trust_remote_code=True
        ).to(device)
        print("✅ Embedding Model loaded successfully.")
        all_results = {}
        for task_name, EvaluatorClass in EVALUATOR_REGISTRY.items():
            print(f"\n--- Evaluating: {task_name} ---")
            try:
                if task_name == "Sentiment Analysis":
                    model = classification_model
                    print("Using classification model for Sentiment Analysis")
                elif task_name in ["Transliteration", "Normalization"]:
                    model = embdding_model
                    print(f"Using base embedding model for {task_name}")
                else: 
                    raise ValueError(f"Unknown task for model selection: {task_name}")
                evaluator: BaseEvaluator = EvaluatorClass()
                result = evaluator.evaluate(model, tokenizer, device=device)
                
                # Extract main metric (must be in every evaluator)
                all_results[task_name] = result["main_metric"]
                print(f"✅ {task_name}: {result['main_metric']:.4f}")
                
            except Exception as e:
                error_msg = f"Failed {task_name}: {str(e)}"
                print(f"❌ {error_msg}")
                all_results[task_name] = None  # or skip

        return EvaluationResult(
            model=model_name,
            revision=revision,
            precision=precision,
            weight_type=weight_type,
            results=all_results
        )

    except Exception as e:
        error_msg = f"Critical failure: {str(e)}"
        print(f"💥 {error_msg}")
        return EvaluationResult(
            model=model_name,
            revision=revision,
            precision=precision,
            weight_type=weight_type,
            results={},
            error=error_msg
        )

def reset_stale_running_eval(eval_entry,root ,file_path ,filename ,timeout_interval=10):
    submission = eval_entry.get("submitted_time")
    try:
        started = datetime.fromisoformat(submission)  # aware datetime
    except Exception as e:
        print("Invalid submitted_time format:", submission, e)

    now_utc = datetime.now(timezone.utc)

    if now_utc - started > timedelta(seconds=timeout_interval):
        print(f"Timeout detected — resetting {eval_entry['model']} to PENDING")
        eval_entry["status"] = EvaluationStatus.PENDING.value
        eval_entry["submitted_time"] = now_utc.isoformat()
        with open(file_path, 'w') as f:
            json.dump(eval_entry, f, indent=2)
        API.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=os.path.join(os.path.basename(root), filename),
            repo_id=QUEUE_REPO,
            repo_type="dataset",
            commit_message=f"Update status to PENDING for {eval_entry['model']} (timeout)",
            token=TOKEN
        )
        return
            
            
            
            
def process_evaluation_queue():
    """
    Processes all pending evaluations in the queue.
    This function acts as a worker that finds a PENDING job, runs it,
    and updates the status on the Hugging Face Hub.
    """
    print("\n=== Starting evaluation queue processing ===")
    print(f"Current time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        
    print(f"Looking for evaluation requests in: {EVAL_REQUESTS_PATH}")
    
    if not os.path.exists(EVAL_REQUESTS_PATH):
        print(f"Evaluation requests path does not exist: {EVAL_REQUESTS_PATH}")
        return
    
    for root, _, files in os.walk(EVAL_REQUESTS_PATH):
        for filename in files:
            if filename.endswith('.json'):
                file_path = os.path.join(root, filename)
                print(f"\nProcessing file: {file_path}")
                
                try:
                    with open(file_path, 'r') as f:
                        eval_entry = json.load(f)
                    
                    status = eval_entry.get('status', '')                    

                    
                    if status == EvaluationStatus.PENDING.value:
                        print(f"Found pending evaluation for model: {eval_entry['model']}")
                        
                        # --- Step 1: Update status to RUNNING locally and on Hub ---
                        eval_entry['status'] = EvaluationStatus.RUNNING.value
                        with open(file_path, 'w') as f:
                            json.dump(eval_entry, f, indent=2)

                        user_name = os.path.basename(root)
                        path_in_repo_queue = os.path.join(user_name, filename)
                        
                        # Upload the updated file to the queue repo to reflect 'RUNNING' status
                        API.upload_file(
                            path_or_fileobj=file_path,
                            path_in_repo=path_in_repo_queue,
                            repo_id=QUEUE_REPO,
                            repo_type="dataset",
                            commit_message=f"Update status to RUNNING for {eval_entry['model']}"
                        )
                        print(f"Updated status to RUNNING in queue: {path_in_repo_queue}")
                        
                        # --- Step 2: Run the evaluation ---
                        print("\n=== Starting evaluation ===")
                        eval_result = evaluate_model(
                            model_name=eval_entry['model'],
                            revision=eval_entry['revision'],
                            precision=eval_entry['precision'],
                            weight_type=eval_entry['weight_type']
                        )
                        
                        for v in eval_result.results.values():
                            if v is None:
                                if eval_result.error is None:
                                    eval_result.error = ""
                                eval_result.error += f"Evaluation failed for {eval_entry['model']}: {v} is None"
                        
                        print("\n=== Evaluation completed ===")
                        
                        # --- Step 3: Update file with final status and results locally ---
                        if eval_result.error:
                            eval_entry['status'] = EvaluationStatus.FAILED.value
                            eval_entry['error'] = eval_result.error
                            print(f"Evaluation failed with error: {eval_result.error}")
                        else:
                            eval_entry['status'] = EvaluationStatus.FINISHED.value
                            eval_entry['results'] = eval_result.results
                            print(f"Evaluation finished successfully. Results: {eval_result.results}")
                        
                        with open(file_path, 'w') as f:
                            json.dump(eval_entry, f, indent=2)
                        
                        # --- Step 4: Upload the final file to the results directory on the Hub ---
                        try:
                            # Use the local file with its final status as the basis for the results file
                            path_in_repo_results = os.path.join(user_name, filename)
                            API.upload_file(
                                path_or_fileobj=file_path,
                                path_in_repo=path_in_repo_results,
                                repo_id=RESULTS_REPO,
                                repo_type="dataset",
                                commit_message=f"Evaluation {'results' if not eval_result.error else 'error'} for {eval_entry['model']}"
                            )
                            print("\nResults uploaded to Hugging Face successfully.")
                            
                        except Exception as upload_error:
                            print(f"Error uploading results: {str(upload_error)}")
                            
                        # --- Step 5: Update the status of the request in the queue to FINISHED/FAILED ---
                        # This keeps a record of all processed jobs in the queue repo.
                        try:
                             API.upload_file(
                                path_or_fileobj=file_path,
                                path_in_repo=path_in_repo_queue,
                                repo_id=QUEUE_REPO,
                                repo_type="dataset",
                                commit_message=f"Final status update for {eval_entry['model']}"
                            )
                             print(f"Final status for {eval_entry['model']} updated in the queue repository.")
                        except Exception as status_update_error:
                            print(f"Error updating status in queue: {str(status_update_error)}")
                    elif status == EvaluationStatus.RUNNING.value:
                        print("Found Running evaluation for model: ", eval_entry['model'])
                        reset_stale_running_eval(eval_entry, root, file_path, filename)
                    else:
                        print(f"Skipping file with status: {status}")
                except Exception as e:
                    print(f"Error processing file {file_path}: {str(e)}")
                    print(f"Full traceback: {traceback.format_exc()}")
                    continue

    print("\n=== Evaluation queue processed. ===")
    print("No more pending jobs found.")
    return