Update app.py

app.py

@@ -1,50 +1,42 @@
 import os
 import torch
-from torch import nn
 from fastapi import FastAPI
 from pydantic import BaseModel, Field
 from comet import load_from_checkpoint
 from huggingface_hub import snapshot_download, HfApi
 
 # ==========================================================
-# ⚙️  Configuração de memória do PyTorch
-# ==========================================================
-os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
-
-# ==========================================================
-# 🚀 Configuração da API
+# 🚀 Configuração da API 
 # ==========================================================
 app = FastAPI(
     title="XCOMET-XXL API",
-    version="1.5.0",
-    description="API para avaliação de traduções usando Unbabel/XCOMET-XXL, "
-                "compatível com campos 'source', 'target' e 'human_translation_ref'."
+    version="2.0.0",
+    description="API para avaliação de traduções usando Unbabel/XCOMET-XXL API "
+               
 )
 
 MODEL_NAME = "Unbabel/XCOMET-XXL"
 HF_TOKEN = os.environ.get("HF_TOKEN")  # defina nas Secrets do Space
-SPACE_REPO_ID = os.environ.get("SPACE_REPO_ID", "nairut/comet-xxl")
+SPACE_REPO_ID = os.environ.get("SPACE_REPO_ID", "nairut/xcomet-xxl")
 
-# ==========================================================
-# 📂 Caminho persistente (150 GB Medium Storage)
-# ==========================================================
-MODEL_DIR = "/data/model"
+# Diretório de cache local (dentro do Space ou ambiente local)
+MODEL_DIR = os.path.join(os.path.dirname(__file__), "model")
 MODEL_CKPT = os.path.join(MODEL_DIR, "checkpoints", "model.ckpt")
 
 
 # ==========================================================
-# ⚙️  Função auxiliar: baixa e persiste o modelo
+# ⚙️ Função auxiliar: baixa e persiste o modelo
 # ==========================================================
 def ensure_model_persisted_once():
     """
-    Faz o download do modelo XCOMET-XXL para /data/model (caso ainda não exista)
+    Faz o download do modelo COMETKiwi-DA-XXL para ./model (caso ainda não exista)
     e tenta commitar essa pasta no próprio Space, para persistência.
     """
     if os.path.exists(MODEL_CKPT):
         print(f"✅ Modelo já existe em {MODEL_CKPT}. Pulando download.")
         return
 
-    print("🔽 Baixando snapshot do modelo para /data/model ...")
+    print("🔽 Baixando snapshot do modelo para ./model ...")
     snapshot_download(
         repo_id=MODEL_NAME,
         token=HF_TOKEN,
@@ -62,7 +54,7 @@ def ensure_model_persisted_once():
             repo_type="space",
             folder_path=MODEL_DIR,
             path_in_repo="model",
-            commit_message="Persistência automática do modelo XCOMET-XXL"
+            commit_message="Persistência automática do modelo COMETKiwi-DA-XXL"
         )
         print("✅ Modelo persistido no Space.")
     except Exception as e:
@@ -71,32 +63,26 @@ def ensure_model_persisted_once():
 
 
 # ==========================================================
-# 📦 Inicialização do modelo (multi-GPU + persistência)
+# ♻️  Inicialização limpa
+# ==========================================================
+# Remove da memória qualquer modelo carregado anteriormente
+if "model" in globals():
+    del model
+    torch.cuda.empty_cache()
+    print("🧹 Modelo anterior removido da memória.")
+
+
+# ==========================================================
+# 📦 Inicialização do modelo
 # ==========================================================
 ensure_model_persisted_once()
 
 print(f"📂 Carregando modelo de {MODEL_CKPT} ...")
-
-# Carrega o modelo (sem map_location, compatível com COMET)
 model = load_from_checkpoint(MODEL_CKPT)
-model.eval()
-
-# Detecta GPUs disponíveis
-num_gpus = torch.cuda.device_count()
-print(f"🎮 GPUs detectadas: {num_gpus}")
-
-# Se houver mais de uma GPU, ativa DataParallel
-if num_gpus > 1:
-    print("⚙️ Ativando DataParallel para usar múltiplas GPUs...")
-    model = nn.DataParallel(model)
-
-# Move o modelo para GPU (caso disponível)
-if torch.cuda.is_available():
-    model.to("cuda")
-    print("✅ Modelo XCOMET-XXL carregado e distribuído nas GPUs.")
-else:
-    print("⚠️ Nenhuma GPU detectada. Rodando em CPU (lento).")
+print("✅ Modelo COMETKiwi-DA-XXL carregado com sucesso!")
 
+USE_GPU = 1 if torch.cuda.is_available() else 0
+print(f"⚙️ GPU detectada: {'sim' if USE_GPU else 'não'}")
 
 
 # ==========================================================
@@ -191,4 +177,4 @@ def score_batch(pairs: list[TranslationPair]):
 # ==========================================================
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)