Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import sys
 import uvicorn
-from fastapi import FastAPI, Query
+from fastapi import FastAPI, Query, HTTPException, BackgroundTasks
 from fastapi.responses import HTMLResponse
 from starlette.middleware.cors import CORSMiddleware
 from datasets import load_dataset, list_datasets
@@ -12,19 +12,23 @@ import psutil
 import asyncio
 import torch
 from tenacity import retry, stop_after_attempt, wait_fixed
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, RepositoryNotFoundError
 from dotenv import load_dotenv
 
+# Cargar variables de entorno
 load_dotenv()
 
+# Obtener el token de Hugging Face
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 if not HUGGINGFACE_TOKEN:
     logger.error("Hugging Face token not found. Please set the HUGGINGFACE_TOKEN environment variable.")
     sys.exit(1)
 
+# Inicializar diccionarios para datasets y ejemplos
 datasets_dict = {}
 example_usage_list = []
 
+# Configuración de caché
 CACHE_DIR = os.path.expanduser("~/.cache/huggingface")
 os.makedirs(CACHE_DIR, exist_ok=True)
 os.environ["HF_HOME"] = CACHE_DIR
@@ -32,6 +36,9 @@ os.environ["HF_TOKEN"] = HUGGINGFACE_TOKEN
 
 pipeline_instance = None  # Solo un pipeline
 
+# Flag para indicar si la inicialización está completa
+initialization_complete = False
+
 def initialize_model():
     global pipeline_instance
     try:
@@ -40,7 +47,7 @@ def initialize_model():
         model = AutoModelForCausalLM.from_pretrained(
             base_model_repo,
             cache_dir=CACHE_DIR,
-            ignore_mismatched_sizes=True  # Añadir este parámetro
+            ignore_mismatched_sizes=True  # Ignorar discrepancias de tamaño
         )
         tokenizer = AutoTokenizer.from_pretrained(base_model_repo, cache_dir=CACHE_DIR)
         if tokenizer.pad_token is None:
@@ -60,7 +67,8 @@ def initialize_model():
 def download_dataset(dataset_name):
     try:
         logger.info(f"Starting download for dataset: {dataset_name}")
-        datasets_dict[dataset_name] = load_dataset(dataset_name, trust_remote_code=True, cache_dir=CACHE_DIR)
+        # Eliminado 'trust_remote_code=True' para evitar el error con ParquetConfig
+        datasets_dict[dataset_name] = load_dataset(dataset_name, cache_dir=CACHE_DIR)
         create_example_usage(dataset_name)
     except Exception as e:
         logger.error(f"Error loading dataset {dataset_name}: {e}", exc_info=True)
@@ -116,6 +124,7 @@ def unify_datasets():
     except Exception as e:
         logger.error(f"Error unifying datasets: {e}", exc_info=True)
 
+# Configuración de concurrencia
 cpu_count = psutil.cpu_count(logical=False) or 1
 memory_available_mb = psutil.virtual_memory().available / (1024 * 1024)
 memory_per_download_mb = 100
@@ -130,24 +139,30 @@ logger.info(f"Using up to {max_concurrent_downloads} concurrent workers for down
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=max_concurrent_downloads)
 
 async def download_and_process_datasets():
-    dataset_names = list_datasets()
-    logger.info(f"Found {len(dataset_names)} datasets to download.")
-    loop = asyncio.get_event_loop()
-    tasks = []
-    for dataset_name in dataset_names:
-        task = loop.run_in_executor(executor, download_dataset, dataset_name)
-        tasks.append(task)
-    await asyncio.gather(*tasks)
-    unify_datasets()
-    upload_model_to_hub()
-
-# Elimina la llamada a asyncio.run(main()) y mueve la inicialización al evento de inicio de FastAPI
+    global initialization_complete
+    try:
+        dataset_names = list_datasets()
+        logger.info(f"Found {len(dataset_names)} datasets to download.")
+        loop = asyncio.get_event_loop()
+        tasks = []
+        for dataset_name in dataset_names:
+            task = loop.run_in_executor(executor, download_dataset, dataset_name)
+            tasks.append(task)
+        await asyncio.gather(*tasks)
+        unify_datasets()
+        upload_model_to_hub()
+        initialization_complete = True
+        logger.info("All initialization tasks completed successfully.")
+    except Exception as e:
+        logger.error(f"Error during dataset processing: {e}", exc_info=True)
 
+# Inicializar FastAPI
 app = FastAPI()
 
+# Configuración de CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # Cambia esto según tus necesidades
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"]
@@ -159,9 +174,19 @@ message_history = []
 async def startup_event():
     logger.info("Application startup initiated.")
     loop = asyncio.get_event_loop()
-    await loop.run_in_executor(None, initialize_model)
-    await download_and_process_datasets()
-    logger.info("Application startup completed.")
+    # Crear una tarea en segundo plano para inicializar el modelo y descargar datasets
+    loop.create_task(run_initialization(loop))
+    logger.info("Background initialization tasks started.")
+
+async def run_initialization(loop):
+    global initialization_complete
+    try:
+        # Inicializar el modelo en un hilo separado
+        await loop.run_in_executor(None, initialize_model)
+        # Descargar y procesar datasets
+        await download_and_process_datasets()
+    except Exception as e:
+        logger.error(f"Error during startup tasks: {e}", exc_info=True)
 
 @app.get('/')
 async def index():
@@ -321,7 +346,12 @@ async def index():
                 userInput.value = '';
 
                 fetch(`/autocomplete?q=${encodeURIComponent(userMessage)}`)
-                .then(response => response.json())
+                .then(response => {
+                    if (response.status === 503) {
+                        return response.json().then(data => { throw new Error(data.detail); });
+                    }
+                    return response.json();
+                })
                 .then(data => {
                     const botMessages = data.result;
                     botMessages.forEach(message => {
@@ -331,6 +361,7 @@ async def index():
                 })
                 .catch(error => {
                     console.error('Error:', error);
+                    appendMessage('bot', error.message || 'An error occurred. Please try again later.');
                 });
             }
 
@@ -349,15 +380,20 @@ async def index():
 
 @app.get('/autocomplete')
 async def autocomplete(q: str = Query(..., title='query')):
-    global message_history
+    global message_history, pipeline_instance, initialization_complete
     message_history.append(('user', q))
+    
+    if not initialization_complete:
+        logger.warning("Model is not initialized yet.")
+        raise HTTPException(status_code=503, detail="Model is not initialized yet. Please try again later.")
+    
     try:
         response = pipeline_instance(q, max_length=50, num_return_sequences=1)[0]['generated_text']
         logger.debug(f"Successfully autocomplete, q:{q}, res:{response}")
         return {"result": [response]}
     except Exception as e:
         logger.error(f"Ignored error in autocomplete: {e}", exc_info=True)
-        return {"result": []}
+        raise HTTPException(status_code=500, detail="An error occurred while processing your request.")
 
 if __name__ == '__main__':
     port = int(os.getenv("PORT", 443))