import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np

def evaluate_perplexity(model_name, revision="main", test_text=None):
    """
    Evaluate perplexity on a fixed piece of text.
    
    Args:
        model_name: Hugging Face model identifier
        revision: Model revision/commit hash
        test_text: Text to evaluate perplexity on (default if None)
    
    Returns:
        float: Perplexity score (lower is better)
    """
    
    # Default test text if none provided
    if test_text is None:
        test_text = """Artificial intelligence has transformed the way we live and work, bringing both opportunities and challenges. 
        From autonomous vehicles to language models that can engage in human-like conversation, AI technologies are becoming increasingly 
        sophisticated. However, with this advancement comes the responsibility to ensure these systems are developed and deployed ethically, 
        with careful consideration for privacy, fairness, and transparency. The future of AI will likely depend on how well we balance innovation 
        with these important social considerations."""
    
    # Load model and tokenizer
    model = AutoModelForCausalLM.from_pretrained(
        model_name, 
        revision=revision,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, revision=revision)
    
    # Tokenize the text
    inputs = tokenizer(test_text, return_tensors="pt")
    
    # Move to same device as model
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    # Calculate loss
    with torch.no_grad():
        outputs = model(**inputs, labels=inputs["input_ids"])
        loss = outputs.loss
    
    # Calculate perplexity
    perplexity = torch.exp(loss).item()
    
    return perplexity

def create_perplexity_result(model_name, revision, precision, perplexity_score):
    """
    Create a result file in the expected format.
    """
    return {
        "config": {
            "model_dtype": f"torch.{precision}",
            "model_name": model_name,
            "model_sha": revision,
        },
        "results": {
            "perplexity": {
                "perplexity": perplexity_score,
            }
        }
    }