Update app.py

app.py

@@ -12,7 +12,7 @@ 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
 
 load_dotenv()
@@ -37,7 +37,11 @@ def initialize_model():
     try:
         logger.info("Initializing the base model and tokenizer.")
         base_model_repo = "meta-llama/Llama-3.2-1B"
-        model = AutoModelForCausalLM.from_pretrained(base_model_repo, cache_dir=CACHE_DIR)
+        model = AutoModelForCausalLM.from_pretrained(
+            base_model_repo,
+            cache_dir=CACHE_DIR,
+            ignore_mismatched_sizes=True  # Añadir este parámetro
+        )
         tokenizer = AutoTokenizer.from_pretrained(base_model_repo, cache_dir=CACHE_DIR)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
@@ -52,8 +56,6 @@ def initialize_model():
         logger.error(f"Error initializing model and tokenizer: {e}", exc_info=True)
         sys.exit(1)
 
-initialize_model()
-
 @retry(stop=stop_after_attempt(3), wait=wait_fixed(5))
 def download_dataset(dataset_name):
     try:
@@ -139,10 +141,7 @@ async def download_and_process_datasets():
     unify_datasets()
     upload_model_to_hub()
 
-async def main():
-    await download_and_process_datasets()
-
-asyncio.run(main())
+# Elimina la llamada a asyncio.run(main()) y mueve la inicialización al evento de inicio de FastAPI
 
 app = FastAPI()
 
@@ -156,6 +155,14 @@ app.add_middleware(
 
 message_history = []
 
+@app.on_event("startup")
+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.")
+
 @app.get('/')
 async def index():
     html_code = """