main.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StaticCache
import re
from typing import List, Tuple
import argparse
import sys
import psutil
import os
import platform
import time

# Add at the very beginning of your script, before any imports
os.environ["TOKENIZERS_PARALLELISM"] = "false"

def get_memory_usage():
    """
    Get the current memory usage of the Python process.
    """
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / (1024 * 1024)  # Convert to MB

def get_gpu_memory_usage(device):
    """
    Get the current GPU memory usage.
    """
    if device.type == "cuda":
        return torch.cuda.memory_allocated() / (1024 * 1024)  # Convert to MB
    elif device.type == "mps":
        return torch.mps.current_allocated_memory() / (1024 * 1024)  # Convert to MB
    return 0

def get_total_gpu_memory(device):
    """
    Get the total GPU memory allocated by the process.
    """
    if device.type == "cuda":
        return torch.cuda.max_memory_allocated() / (1024 * 1024)  # Convert to MB
    elif device.type == "mps":
        return torch.mps.driver_allocated_memory() / (1024 * 1024)  # Convert to MB
    return 0

def get_available_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")
        
def load_model_and_tokenizer(model_name, device):
    print("Loading tokenizer and model...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        if device.type == "cuda":
            # For CUDA, we can use device_map="auto"
            model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
        else:
            # For CPU and MPS, we need to explicitly move the model
            model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
        
        return tokenizer, model
    except Exception as e:
        print(f"Error loading model or tokenizer: {e}")
        sys.exit(1)
        
def extract_answer_span(generated_text: str) -> str:
    """
    Extracts the answer span from the generated text.
    Assumes that the answer starts after the last occurrence of 'Answer:'.

    Args:
        generated_text (str): The text generated by the model.

    Returns:
        str: The extracted answer span.
    """
    match = re.search(r'Answer:\s*(.*)', generated_text, re.IGNORECASE | re.DOTALL)
    if match:
        return match.group(1).strip()
    else:
        return ""

def calculate_cot_score(model, tokenizer, generated_text: str, device: torch.device) -> float:
    """
    Calculates the Chain-of-Thought (CoT) score for the generated answer.

    Args:
        model: The language model.
        tokenizer: The tokenizer corresponding to the model.
        generated_text (str): The full text generated by the model.
        device: The device to run the model on.

    Returns:
        float: The CoT score.
    """
    # Extract the answer span
    answer = extract_answer_span(generated_text)
    if not answer:
        return 0.0

    # Tokenize the generated text
    inputs = tokenizer(generated_text, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=False)
        logits = outputs.logits  # Shape: [1, seq_len, vocab_size]

    # Get the token ids for the generated text
    generated_ids = inputs['input_ids'][0]

    # Find the starting token index of the answer
    answer_start = generated_text.lower().rfind('answer:')
    if answer_start == -1:
        return 0.0

    # Calculate the number of tokens before the answer
    prefix = generated_text[:answer_start]
    prefix_ids = tokenizer(prefix, return_tensors="pt").input_ids.to(device)
    prefix_length = prefix_ids.size(1)

    # Get the token ids for the answer
    answer_text = generated_text[answer_start:]
    answer_ids = tokenizer(answer_text, return_tensors="pt").input_ids.to(device)[0]

    # Iterate over the answer tokens and calculate p(x1t) - p(x2t)
    cot_scores = []
    for i in range(1, len(answer_ids)):
        token_position = prefix_length + i - 1  # Adjust for 0-based index
        if token_position >= logits.size(1) - 1:
            break  # Avoid index out of range

        logits_at_t = logits[0, token_position, :]
        probs = torch.softmax(logits_at_t, dim=-1)
        top_probs, top_indices = torch.topk(probs, 2)
        if len(top_probs) < 2:
            continue
        delta = top_probs[0].item() - top_probs[1].item()
        cot_scores.append(delta)

    if not cot_scores:
        return 0.0

    # Calculate the average delta
    average_delta = sum(cot_scores) / len(cot_scores)
    return average_delta

def cot_decoding(model, tokenizer, prompt_messages: List[dict], k: int = 10, aggregation: str = 'max', device: torch.device = 'cpu', use_sequential: bool = False) -> Tuple[str, float]:
    # Prepare the tokenized prompt
    tokenized_chat = tokenizer.apply_chat_template(
        prompt_messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True
    ).to(device)

    input_ids = tokenized_chat['input_ids']
    attention_mask = tokenized_chat.get('attention_mask', None)

    # Set pad_token_id if it's not already set
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # Initialize StaticCache for the prompt 
    static_cache = StaticCache(
        config=model.config,
        batch_size=1,
        max_cache_len=input_ids.shape[1] + 256,  # Prompt length + max new tokens
        device=device,
        dtype=model.dtype
    )

    # Get logits for the next token
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, past_key_values=static_cache, use_cache=True)
        logits = outputs.logits
        next_token_logits = logits[0, -1, :]

    # Get top-k tokens for the first generation step
    topk_probs, topk_indices = torch.topk(torch.softmax(next_token_logits, dim=-1), k)

    generated_texts = []
    cot_scores = []
    
    prompt_length = input_ids.shape[1]

    # Measure the time taken for all branches
    start_time = time.time()

    if use_sequential:
        for i in range(k):
            token_id = topk_indices[i].unsqueeze(0).unsqueeze(0)
            current_input_ids = torch.cat([input_ids, token_id], dim=-1)
            current_attention_mask = torch.cat([attention_mask, torch.ones((1, 1), device=device, dtype=torch.long)], dim=-1)
            
            with torch.no_grad():
                generated_output = model.generate(
                    input_ids=current_input_ids,
                    attention_mask=current_attention_mask,
                    max_new_tokens=256,
                    do_sample=False,
                    top_p=None,
                    top_k=None,
                    temperature=None,
                    num_return_sequences=1,
                    use_cache=True,
                    pad_token_id=tokenizer.pad_token_id,
                    eos_token_id=tokenizer.eos_token_id,
                    cache_position=torch.arange(input_ids.shape[1], current_input_ids.shape[1], device=device)
                )
            
            # Extract only the newly generated part
            generated_text = tokenizer.decode(generated_output[0][prompt_length:], skip_special_tokens=True).strip()
            generated_texts.append(generated_text)
            
            score = calculate_cot_score(model, tokenizer, generated_text, device)
            cot_scores.append(score)
            
            print(f"Path {i+1}:\n{'='*25}\n{generated_text}\nAnswer span: {extract_answer_span(generated_text)}\nCoT Score: {score}\n\n")
            
            # Empty torch cache depending on the device
            if device.type == "cuda":
                torch.cuda.empty_cache()
            elif device.type == "mps":
                torch.mps.empty_cache()

    else:
        # Prepare k different input sequences, each starting with a different top-k token
        branches_input_ids = torch.cat([input_ids.repeat(k, 1), topk_indices.unsqueeze(1)], dim=1)
        branches_attention_mask = torch.cat([attention_mask.repeat(k, 1), torch.ones((k, 1), device=device, dtype=torch.long)], dim=1)

        # Generate the rest of the sequence greedily for each branch
        with torch.no_grad():
            generated_outputs = model.generate(
                input_ids=branches_input_ids,
                attention_mask=branches_attention_mask,
                max_new_tokens=256,
                do_sample=False,
                top_p=None,
                top_k=None,
                temperature=None,
                num_return_sequences=1,
                use_cache=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
                cache_position=torch.arange(input_ids.shape[1], branches_input_ids.shape[1], device=device).repeat(k, 1)
            )

        for i, generated_ids in enumerate(generated_outputs):
            # Extract only the newly generated part
            generated_text = tokenizer.decode(generated_ids[prompt_length:], skip_special_tokens=True).strip()
            generated_texts.append(generated_text)
            
            score = calculate_cot_score(model, tokenizer, generated_text, device)
            cot_scores.append(score)
            
            print(f"Path {i+1}:\n{'='*25}\n{generated_text}\nAnswer span: {extract_answer_span(generated_text)}\nCoT Score: {score}\n\n")

    end_time = time.time()
    print(f"Time taken for all branches: {end_time - start_time:.2f} seconds")

    # Aggregation logic remains the same
    if aggregation == 'max':
        max_score = max(cot_scores)
        max_index = cot_scores.index(max_score)
        selected_answer = extract_answer_span(generated_texts[max_index])
        return selected_answer, max_score
    elif aggregation == 'sum':
        answer_score_map = {}
        for answer, score in zip([extract_answer_span(text) for text in generated_texts], cot_scores):
            answer_score_map[answer] = answer_score_map.get(answer, 0) + score

        print("Aggregation Mapping (Answer: Total CoT Score):")
        for ans, total_score in answer_score_map.items():
            print(f"{ans}: {total_score}")

        selected_answer = max(answer_score_map, key=answer_score_map.get)
        sum_scores = answer_score_map[selected_answer]
        return selected_answer, sum_scores
    else:
        raise ValueError("Invalid aggregation method. Choose 'max' or 'sum'.")

def main():
    parser = argparse.ArgumentParser(description="Demonstrate Chain-of-Thought (CoT) Decoding using Transformers")
    parser.add_argument('--model_name', type=str, default="meta-llama/Llama-3.2-1B-Instruct", help='Model checkpoint name')
    parser.add_argument('--question', type=str, default="Sally has two brothers, Sam and Joe. Sam has one sister. How many sisters does Joe have? Think step by step. You MUST end your reply with Answer:, FOLLOWED BY A SINGLE NUMBER.", help='Question to ask the model')
    parser.add_argument('--k', type=int, default=10, help='Number of decoding branches')
    parser.add_argument('--aggregation', type=str, default='max', choices=['max', 'sum'], help='Aggregation method for CoT scores')
    parser.add_argument('--device', type=str, default='cuda', choices=['cuda', 'cpu', 'mps'], help='Device to run the model on')
    parser.add_argument('--use_sequential', action='store_true', help='Use sequential processing for low RAM situations')
    parser.add_argument('--system_prompt', type=str, default=False,help='Optional system prompt for the model')
    args = parser.parse_args()

    max_memory_usage = 0
    max_gpu_memory_usage = 0
    
    # Determine the available device
    device = get_available_device()
    print(f"Using device: {device}")

    # Load tokenizer and model
    tokenizer, model = load_model_and_tokenizer(args.model_name, device)
    max_memory_usage = max(max_memory_usage, get_memory_usage())
    max_gpu_memory_usage = max(max_gpu_memory_usage, get_gpu_memory_usage(device))

    model.eval()

    # Prepare the prompt messages
    messages = [
        {
            "role": "user",
            "content": args.question,
        }
    ]
    
    if args.system_prompt:
        messages.insert(0, {
            "role": "system",
            "content": args.system_prompt
        })
        
    # Perform CoT decoding
    print(f"Performing CoT decoding with k={args.k}, aggregation='{args.aggregation}', and use_sequential={args.use_sequential}...\n")
    answer, score = cot_decoding(model, tokenizer, messages, k=args.k, aggregation=args.aggregation, device=device, use_sequential=args.use_sequential)
    max_memory_usage = max(max_memory_usage, get_memory_usage())
    max_gpu_memory_usage = max(max_gpu_memory_usage, get_gpu_memory_usage(device))

    # Display the result
    print("\n=== CoT Decoding Result ===")
    print(f"Question: {args.question}")
    print(f"Answer: {answer}")
    print(f"CoT Score: {score}")
    print(f"\nMaximum RAM usage: {max_memory_usage:.2f} MB")
    
    if device.type in ['cuda', 'mps']:
        print(f"Maximum GPU memory usage: {max_gpu_memory_usage:.2f} MB")
        print(f"Total GPU memory allocated: {get_total_gpu_memory(device):.2f} MB")
        
    if device.type == 'cuda':
        print(f"Maximum total VRAM allocated: {torch.cuda.max_memory_allocated() / (1024 * 1024):.2f} MB")
        print(f"Maximum total VRAM reserved: {torch.cuda.max_memory_reserved() / (1024 * 1024):.2f} MB")
    elif device.type == 'mps':
        print("Note: On MPS devices, 'Maximum GPU memory usage' shows the peak tensor memory,")
        print("while 'Total GPU memory allocated' shows the total memory allocated by the Metal driver.")

if __name__ == "__main__":
    main()