import gradio as gr
import subprocess
import spaces
import os
import json

# Ensure the results directory exists
os.makedirs('./eval_results', exist_ok=True)

def run_evaluation():
    # Fetch the token from Space Secrets
    hf_token = os.environ.get("HF_TOKEN", "")
    
    if not hf_token:
        yield "⚠️ Warning: HF_TOKEN secret not found. Attempting unauthenticated access..."
    else:
        yield "🔐 Authenticating terminal session with HF_TOKEN..."
        subprocess.run(["huggingface-cli", "login", "--token", hf_token], capture_output=True)
        yield "✅ Authentication successful."

    tasks = "mmlu,gsm8k,winogrande"
    output_file = "./eval_results/results.json"
    
    # CHANGED: Removed load_in_4bit=True, added device=cuda
    command = [
        "lm_eval",
        "--model", "hf",
        "--model_args", "pretrained=Girikannan/sarvam-eval,device_map=auto,trust_remote_code=True",
        "--device", "cuda",
        "--tasks", tasks,
        "--batch_size", "1",
        "--output_path", output_file
    ]
    
    yield "🚀 Initializing evaluation on GPU... (The model is already downloaded and cached, so this will boot up fast!)"
    
    try:
        process = subprocess.run(command, capture_output=True, text=True)
        
        if process.returncode == 0:
            with open(output_file, 'r') as f:
                results = json.load(f)
            yield f"🎯 Evaluation Complete!\n\nResults: {json.dumps(results, indent=2)}"
        else:
            yield f"❌ Evaluation failed with error:\n{process.stderr}"
            
    except Exception as e:
        yield f"🚨 An unexpected error occurred: {str(e)}"
        

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Sarvam-30b Compression Evaluator")
    gr.Markdown("Click the button below to run the evaluation suite.")
    
    eval_button = gr.Button("Start Full Evaluation")
    output_text = gr.Textbox(label="Evaluation Log", lines=20)
    
    run_eval_gpu = spaces.GPU(duration=3600)(run_evaluation)
    eval_button.click(run_eval_gpu, inputs=None, outputs=output_text)

if __name__ == "__main__":
    demo.launch()