Update app.py

app.py

@@ -4,7 +4,6 @@ from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from transformers import (
-    AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
@@ -15,6 +14,7 @@ import asyncio
 from io import BytesIO
 from botocore.exceptions import NoCredentialsError
 import boto3
+from huggingface_hub import snapshot_download
 
 # Diccionario global para almacenar los tokens y configuraciones de los modelos
 token_dict = {}
@@ -64,12 +64,16 @@ class S3ModelLoader:
         
         s3_uri = self._get_s3_uri(model_name)
         try:
-            model = AutoModelForCausalLM.from_pretrained(s3_uri, local_files_only=True)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri, local_files_only=True)
+            # Descargamos el modelo y el tokenizer desde Hugging Face directamente a S3
+            model_path = snapshot_download(model_name, token=HUGGINGFACE_HUB_TOKEN)
+            
+            model = AutoModelForCausalLM.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_path)
 
             if tokenizer.eos_token_id is None:
                 tokenizer.eos_token_id = tokenizer.pad_token_id
 
+            # Guardamos en el diccionario global
             token_dict[model_name] = {
                 "model": model,
                 "tokenizer": tokenizer,
@@ -77,6 +81,10 @@ class S3ModelLoader:
                 "eos_token_id": tokenizer.eos_token_id
             }
 
+            # Subimos los modelos al S3 si es necesario
+            self.s3_client.upload_file(model_path, self.bucket_name, f'{model_name}/model')
+            self.s3_client.upload_file(f'{model_path}/tokenizer', self.bucket_name, f'{model_name}/tokenizer')
+
             return token_dict[model_name]
         except NoCredentialsError:
             raise HTTPException(status_code=500, detail="AWS credentials not found.")