Update app.py

app.py

@@ -2,12 +2,14 @@ import gradio as gr
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
-
-
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 # Nombre del modelo en el repositorio de Hugging Face
 model_name = "mkjaramillo/cancer2"
-model = torch.load("model_name")
+
+# Cargar el tokenizer y el modelo desde el repositorio de Hugging Face
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
 # Transformación de la imagen
 image_transform = transforms.Compose([
@@ -16,7 +18,6 @@ image_transform = transforms.Compose([
    
 ])
 
-# Función para realizar la predicción
 def classify_image(image):
     # Cargar la imagen
     image = Image.fromarray(image)
@@ -28,10 +29,10 @@ def classify_image(image):
     outputs = model(image)
     
     # Obtener las predicciones
-    predictions = torch.argmax(outputs, dim=1)
+    predictions = torch.argmax(outputs.logits, dim=1)
     
     # Obtener la etiqueta de la predicción
-    label = predictions.item()
+    label = tokenizer.decode(predictions.item())
 
     # Retornar la etiqueta de la predicción
     return label