run_epicache.py

# For licensing see accompanying LICENSE file.
# Copyright (C) 2025 Apple Inc. All Rights Reserved.

import json
import os
from collections import defaultdict
from tqdm import tqdm
import torch
import numpy as np

from args import args
from data import load_dataset_all
from model import LongConvQAModel
from utils import f1_score
from utils.cluster import ClusterManager
from utils.func import TimeStamp

if __name__ == "__main__":
    
    # ========================================================= #
    # Load Model
    # ========================================================= #

    model = LongConvQAModel(args.model, dtype=args.dtype, evict_level=args.level, scoring_method=args.scoring_method)

    # Initialize ClusterManager
    cluster_manager = ClusterManager(
        embedding_type=args.embedding_type,
        n_clusters=args.n_cluster,
        conv_window=args.conv_window,
        medoid_number=args.n_medoid,
        verbose=args.verbose
    )

    if args.scoring_method == "clustering":
        # ========================================================= #
        # Offline Clustering (Run Once)
        # ========================================================= #
        print(">>> Running Offline Clustering...")

        if args.data == "longmemeval":
            data_path = f"data/longmemeval/custom_lme_{args.target_length}_50.json"
            with open(data_path, 'r') as f:
                dataset = json.load(f)
            dataset = dataset['conversations']
        else:
            data_path = "./data/locomo/locomo10.json" if args.data == "locomo" else "./data/realtalk/realtalk10.json"
            with open(data_path, 'r') as f:
                dataset = json.load(f)
        
        for conv_idx in tqdm(range(len(dataset)), desc=f"Conversation Clustering...", disable=not args.verbose, leave=False):
            
            conversation = dataset[conv_idx]
            
            # ============================================================== #
            # Embed and cluster conversation windows
            # ============================================================== #
            if args.data == "longmemeval":
                conversation_data = conversation
                conversation_windows = cluster_manager.extract_conversation_windows_longmemeval(conversation_data)
            else:
                conversation_data = conversation['conversation'] if args.data == "locomo" else conversation
                conversation_windows = cluster_manager.extract_conversation_windows(conversation_data)

            embedded_windows = cluster_manager.embed_conversations(conversation_windows, model=model)
            clustering_results = cluster_manager.cluster_conversations(embedded_windows)
            
            # Save results for this conversation
            conversation_result = {
                'conv_idx': conv_idx,
                'embedded_windows': embedded_windows,
                'clustering_results': clustering_results,
                'n_windows': len(embedded_windows)
            }
            cluster_manager.all_conversation_results.append(conversation_result)
            
            for cluster_id in range(cluster_manager.n_clusters):
                cluster_info = clustering_results['cluster_results'][cluster_id]
                
    
    # ========================================================= #
    # Online Evaluation (LLM Inference)
    # ========================================================= #
    dataset, _ = load_dataset_all(args.data, model.tokenizer, target_length=args.target_length)  # list of data

    # Store results for each question
    all_results = {}
    type_scores = defaultdict(list)
    total_scores = []
    
    print(">>> Running Online Evaluation...")
    for data_idx, data in enumerate(dataset):
        
        # ========================================================= #
        # Prefilling (Build Episodic KV Cache)
        # ========================================================= #
        kvs = []
        for cluster_idx in tqdm(range(cluster_manager.n_clusters), desc=f"Prefilling KV Cache for episodes for conv {data_idx}", disable=not args.verbose, leave=False):
            # Get clustering results for this conversation
            clustering_results = cluster_manager.all_conversation_results[data_idx]['clustering_results']
            
            # Create cluster prompt
            combined_text = cluster_manager.make_cluster_prompt(clustering_results['cluster_results'][cluster_idx]['windows'])
            
            # Prefill
            model.cluster_ids = model.encode(combined_text)
            ctx_ids = model.encode(data['context'])
            kv = model.prefill_memory_constrained(ctx_ids, prefill_chunk_size=args.prefill_chunk_size, \
                score_path=args.score_path, kv_budget=args.kv_budget, power=args.power)

            if args.verbose and cluster_idx == 0 and data_idx == 0:
                kv_budget = kv._budget()
                print(f">>> Avg Budget: {kv_budget.mean().item()} | Max Budget: {kv_budget.max().item()} | Min Budget: {kv_budget.min().item()} | Memory Usage: {kv._mem()} GB")
            
            # Move kv cache to CPU for memory efficiency
            kv.to_cpu()
            kvs.append(kv)
            model._init_kv()

        torch.cuda.empty_cache()
        # ========================================================= #
        # Decoding Answer (Query-Episode Matching based)
        # ========================================================= #
        clustering_results = cluster_manager.all_conversation_results[data_idx]['clustering_results']
        centroids = [clustering_results['cluster_results'][i]['centroid'] for i in range(cluster_manager.n_clusters)]
        
        for question_idx in tqdm(
            range(len(data['question'])),
            desc=f"Question answering for episodes for conv {data_idx}",
            disable=not args.verbose,
            leave=False 
        ):
            question = data['question'][question_idx]
            answer = data['answers'][question_idx]
            question_type = data['task_types'][question_idx]
            
            # 1. Query Embedding
            question_embedding = cluster_manager.embed_question(question, model=model)
            
            # 2. Cluster matching
            similarities = []
            for centroid in centroids:
                similarity = float(np.dot(question_embedding, centroid))
                similarities.append(similarity)
            best_cluster_idx = similarities.index(max(similarities))
            
            # 3. KV retrieval
            selected_kv = kvs[best_cluster_idx]
            selected_kv.to_gpu(model.device)
            
            if args.data == "realtalk" and args.model == "Qwen/Qwen2.5-3B-Instruct":
                question = "\n".join([q for i, q in enumerate(question.split("\n")) if i != 1])

            # 4. Generation
            pred, num_generated_tokens = model.generate(model.apply_template(question), kv=selected_kv)
            score = f1_score(pred, answer)

            # Store result for this question
            result_key = f"conv_{data_idx}_cluster_{best_cluster_idx}_q_{question_idx}"
            all_results[result_key] = {
                'conv_idx': data_idx,
                'cluster_idx': best_cluster_idx,
                'question_idx': question_idx,
                'question': question,
                'prediction': pred,
                'ground_truth': answer,
                'f1_score': score,
                'type': question_type,
                'similarity_to_cluster': max(similarities)
            }
            
            # Track scores for type-wise and overall averages
            type_scores[question_type].append(score)
            total_scores.append(score)
        
    # Calculate average scores
    type_averages = {}
    for type_name, scores in type_scores.items():
        type_averages[type_name] = {
            'average_f1': sum(scores) / len(scores),
            'count': len(scores)
        }
    
    overall_average = sum(total_scores) / len(total_scores) if total_scores else 0
    
    # Combine all results
    final_results = {
        'individual_results': all_results,
        'type_averages': type_averages,
        'overall_average': {
            'average_f1': overall_average,
            'total_count': len(total_scores)
        }
    }
    
    # Save to JSON file
    output_file_name = f"cluster_{args.model.split('/')[-1]}_{args.data}.json"
    output_dir = f"results/{args.exp_name}"
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    output_file = os.path.join(output_dir, output_file_name)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(final_results, f, indent=2, ensure_ascii=False)
    
    print(f"\n=== Evaluation Results ===")
    print(f"Overall Average F1: {overall_average:.4f} ({len(total_scores)} questions)")
    print(f"Type-wise averages:")
    for type_name, stats in type_averages.items():
        print(f"  {type_name}: {stats['average_f1']:.4f} ({stats['count']} questions)")
    print(f"Results saved to: {output_file}")
    
    print("Finished.")