import torch
import torch.nn as nn
from torchvision import models
import gradio as gr
import os
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from safetensors.torch import load_model
from datasets import load_dataset
from models import FromZero, PreTrained

from utils import cargar_etiquetas
from dataset import cargar_dataset
from evaluation import evaluate_interface

# Cargar etiquetas
etiquetas, num_clases, codigo = cargar_etiquetas()

# Cargar dataset
test_dataloader = cargar_dataset(codigo)


def interface_wrapper(model_file, model_type):
    return evaluate_interface(model_file, model_type, num_clases, test_dataloader)


# Interfaz de Gradio
demo = gr.Interface(
    fn=interface_wrapper,
    inputs=[
        gr.File(label="Archivo del modelo (.safetensor)"),
        gr.Radio(["tarea_7", "tarea_8"], label="Tipo de modelo", value="tarea_7"),
    ],
    outputs=gr.Textbox(label="Resultado", lines=1),
    title="Evaluador de Tareas 7 y 8",
    description="Carga un archivo .safetensor de la tarea 7 o 8 y evalúa su precisión en el conjunto de datos de evaluación.",
)

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