import gradio as gr
import spaces
from openmodellab.genome.model_loader import load_model
from openmodellab.genome.analyzer import analyze_model
from openmodellab.benchmark import analyze_benchmark


@spaces.GPU
def analyze(model_name):
    try:
        model, tokenizer = load_model(model_name)

        genome = analyze_model(
            model_name,
            model,
            tokenizer
        )

        benchmark = analyze_benchmark(
            model,
            tokenizer
        )

        return f"""
## OpenModelLab Results

**Model:** `{model_name}`

### Architecture

- Parameters: {genome["model"].get("parameter_count", "N/A") / 1_000_000:.2f}M
- Architecture: {genome["model"].get("architecture", "N/A")}

### Hardware

- Device: {benchmark["hardware"].get("device")}
- GPU: {benchmark["hardware"].get("gpu_name", "N/A")}
- CUDA: {benchmark["hardware"].get("cuda_version", "N/A")}

### Benchmark

- Latency: **{benchmark["latency"]["average_ms"]} ms**
- Throughput: **{benchmark["throughput"]["samples_per_second"]} samples/sec**
- Peak GPU memory: **{benchmark["memory"].get("gpu_peak_allocated_mb", "N/A")} MB**

### Batch Scaling

{benchmark["batch_scaling"]["results"]}

[Whitepaper / DOI](https://doi.org/10.5281/zenodo.21861654)
"""

    except Exception as e:
        return f"Error: {type(e).__name__}: {e}"


demo = gr.Interface(
    fn=analyze,
    inputs=gr.Textbox(
        label="Hugging Face Model",
        value="sentence-transformers/all-MiniLM-L6-v2"
    ),
    outputs=gr.Markdown(),
    title="OpenModelLab",
    description="AI model genome extraction and runtime benchmarking."
)

demo.launch()