Spaces:
Sleeping
Sleeping
Hjgugugjhuhjggg
commited on
Delete app.py.bkk
Browse files- app.py.bkk +0 -203
app.py.bkk
DELETED
@@ -1,203 +0,0 @@
|
|
1 |
-
import os
|
2 |
-
import json
|
3 |
-
import logging
|
4 |
-
import boto3
|
5 |
-
from fastapi import FastAPI, HTTPException, Query
|
6 |
-
from fastapi.responses import JSONResponse
|
7 |
-
from transformers import AutoModelForCausalLM, AutoTokenizer
|
8 |
-
from huggingface_hub import hf_hub_download
|
9 |
-
import asyncio
|
10 |
-
|
11 |
-
logger = logging.getLogger(__name__)
|
12 |
-
logger.setLevel(logging.INFO)
|
13 |
-
console_handler = logging.StreamHandler()
|
14 |
-
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
|
15 |
-
console_handler.setFormatter(formatter)
|
16 |
-
logger.addHandler(console_handler)
|
17 |
-
|
18 |
-
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
|
19 |
-
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
|
20 |
-
AWS_REGION = os.getenv("AWS_REGION")
|
21 |
-
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
|
22 |
-
HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
|
23 |
-
|
24 |
-
MAX_TOKENS = 1024
|
25 |
-
|
26 |
-
s3_client = boto3.client(
|
27 |
-
's3',
|
28 |
-
aws_access_key_id=AWS_ACCESS_KEY_ID,
|
29 |
-
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
|
30 |
-
region_name=AWS_REGION
|
31 |
-
)
|
32 |
-
|
33 |
-
app = FastAPI()
|
34 |
-
|
35 |
-
class S3DirectStream:
|
36 |
-
def __init__(self, bucket_name):
|
37 |
-
self.s3_client = boto3.client(
|
38 |
-
's3',
|
39 |
-
aws_access_key_id=AWS_ACCESS_KEY_ID,
|
40 |
-
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
|
41 |
-
region_name=AWS_REGION
|
42 |
-
)
|
43 |
-
self.bucket_name = bucket_name
|
44 |
-
|
45 |
-
async def stream_from_s3(self, key):
|
46 |
-
loop = asyncio.get_event_loop()
|
47 |
-
return await loop.run_in_executor(None, self._stream_from_s3, key)
|
48 |
-
|
49 |
-
def _stream_from_s3(self, key):
|
50 |
-
try:
|
51 |
-
response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
|
52 |
-
return response['Body']
|
53 |
-
except self.s3_client.exceptions.NoSuchKey:
|
54 |
-
raise HTTPException(status_code=404, detail=f"El archivo {key} no existe en el bucket S3.")
|
55 |
-
except Exception as e:
|
56 |
-
raise HTTPException(status_code=500, detail=f"Error al descargar {key} desde S3: {str(e)}")
|
57 |
-
|
58 |
-
async def get_model_file_parts(self, model_name):
|
59 |
-
loop = asyncio.get_event_loop()
|
60 |
-
return await loop.run_in_executor(None, self._get_model_file_parts, model_name)
|
61 |
-
|
62 |
-
def _get_model_file_parts(self, model_name):
|
63 |
-
try:
|
64 |
-
model_name = model_name.replace("/", "-").lower()
|
65 |
-
files = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=model_name)
|
66 |
-
model_files = [obj['Key'] for obj in files.get('Contents', []) if model_name in obj['Key']]
|
67 |
-
return model_files
|
68 |
-
except Exception as e:
|
69 |
-
raise HTTPException(status_code=500, detail=f"Error al obtener archivos del modelo {model_name} desde S3: {e}")
|
70 |
-
|
71 |
-
async def load_model_from_s3(self, model_name):
|
72 |
-
try:
|
73 |
-
model_name = model_name.replace("/", "-").lower()
|
74 |
-
model_files = await self.get_model_file_parts(model_name)
|
75 |
-
|
76 |
-
if not model_files:
|
77 |
-
await self.download_and_upload_to_s3(model_name)
|
78 |
-
|
79 |
-
config_stream = await self.stream_from_s3(f"{model_name}/config.json")
|
80 |
-
config_data = config_stream.read()
|
81 |
-
|
82 |
-
if not config_data:
|
83 |
-
raise HTTPException(status_code=500, detail=f"El archivo de configuración {model_name}/config.json está vacío o no se pudo leer.")
|
84 |
-
|
85 |
-
config_text = config_data.decode("utf-8")
|
86 |
-
config_json = json.loads(config_text)
|
87 |
-
|
88 |
-
model = AutoModelForCausalLM.from_pretrained(f"s3://{self.bucket_name}/{model_name}", config=config_json, from_tf=False)
|
89 |
-
return model
|
90 |
-
|
91 |
-
except HTTPException as e:
|
92 |
-
raise e
|
93 |
-
except Exception as e:
|
94 |
-
raise HTTPException(status_code=500, detail=f"Error al cargar el modelo desde S3: {e}")
|
95 |
-
|
96 |
-
async def load_tokenizer_from_s3(self, model_name):
|
97 |
-
try:
|
98 |
-
model_name = model_name.replace("/", "-").lower()
|
99 |
-
tokenizer_stream = await self.stream_from_s3(f"{model_name}/tokenizer.json")
|
100 |
-
tokenizer_data = tokenizer_stream.read().decode("utf-8")
|
101 |
-
|
102 |
-
tokenizer = AutoTokenizer.from_pretrained(f"s3://{self.bucket_name}/{model_name}")
|
103 |
-
return tokenizer
|
104 |
-
except Exception as e:
|
105 |
-
raise HTTPException(status_code=500, detail=f"Error al cargar el tokenizer desde S3: {e}")
|
106 |
-
|
107 |
-
async def create_s3_folders(self, s3_key):
|
108 |
-
try:
|
109 |
-
folder_keys = s3_key.split('-')
|
110 |
-
for i in range(1, len(folder_keys)):
|
111 |
-
folder_key = '-'.join(folder_keys[:i]) + '/'
|
112 |
-
if not await self.file_exists_in_s3(folder_key):
|
113 |
-
logger.info(f"Creando carpeta en S3: {folder_key}")
|
114 |
-
self.s3_client.put_object(Bucket=self.bucket_name, Key=folder_key, Body='')
|
115 |
-
|
116 |
-
except Exception as e:
|
117 |
-
raise HTTPException(status_code=500, detail=f"Error al crear carpetas en S3: {e}")
|
118 |
-
|
119 |
-
async def file_exists_in_s3(self, s3_key):
|
120 |
-
try:
|
121 |
-
self.s3_client.head_object(Bucket=self.bucket_name, Key=s3_key)
|
122 |
-
return True
|
123 |
-
except self.s3_client.exceptions.ClientError:
|
124 |
-
return False
|
125 |
-
|
126 |
-
async def download_and_upload_to_s3(self, model_name, force_download=False):
|
127 |
-
try:
|
128 |
-
if force_download:
|
129 |
-
logger.info(f"Forzando la descarga del modelo {model_name} y la carga a S3.")
|
130 |
-
|
131 |
-
model_name = model_name.replace("/", "-").lower()
|
132 |
-
|
133 |
-
if not await self.file_exists_in_s3(f"{model_name}/config.json") or not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
|
134 |
-
config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
|
135 |
-
tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, force_download=force_download)
|
136 |
-
|
137 |
-
await self.create_s3_folders(f"{model_name}/")
|
138 |
-
|
139 |
-
if not await self.file_exists_in_s3(f"{model_name}/config.json"):
|
140 |
-
with open(config_file, "rb") as file:
|
141 |
-
self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
|
142 |
-
|
143 |
-
if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
|
144 |
-
with open(tokenizer_file, "rb") as file:
|
145 |
-
self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
|
146 |
-
else:
|
147 |
-
logger.info(f"Los archivos del modelo {model_name} ya existen en S3. No es necesario descargarlos de nuevo.")
|
148 |
-
|
149 |
-
except Exception as e:
|
150 |
-
raise HTTPException(status_code=500, detail=f"Error al descargar o cargar archivos desde Hugging Face a S3: {e}")
|
151 |
-
|
152 |
-
async def resume_download(self, model_name):
|
153 |
-
try:
|
154 |
-
logger.info(f"Reanudando la descarga del modelo {model_name} desde Hugging Face.")
|
155 |
-
config_file = hf_hub_download(repo_id=model_name, filename="config.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
|
156 |
-
tokenizer_file = hf_hub_download(repo_id=model_name, filename="tokenizer.json", token=HUGGINGFACE_HUB_TOKEN, resume_download=True)
|
157 |
-
|
158 |
-
if not await self.file_exists_in_s3(f"{model_name}/config.json"):
|
159 |
-
with open(config_file, "rb") as file:
|
160 |
-
self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/config.json", Body=file)
|
161 |
-
|
162 |
-
if not await self.file_exists_in_s3(f"{model_name}/tokenizer.json"):
|
163 |
-
with open(tokenizer_file, "rb") as file:
|
164 |
-
self.s3_client.put_object(Bucket=self.bucket_name, Key=f"{model_name}/tokenizer.json", Body=file)
|
165 |
-
|
166 |
-
except Exception as e:
|
167 |
-
raise HTTPException(status_code=500, detail=f"Error al reanudar la descarga o cargar archivos desde Hugging Face a S3: {e}")
|
168 |
-
|
169 |
-
def split_text_by_tokens(text, tokenizer, max_tokens=MAX_TOKENS):
|
170 |
-
tokens = tokenizer.encode(text)
|
171 |
-
chunks = []
|
172 |
-
for i in range(0, len(tokens), max_tokens):
|
173 |
-
chunk = tokens[i:i+max_tokens]
|
174 |
-
chunks.append(tokenizer.decode(chunk))
|
175 |
-
return chunks
|
176 |
-
|
177 |
-
def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
|
178 |
-
generated_text = ""
|
179 |
-
while len(input_text) > 0:
|
180 |
-
chunks = split_text_by_tokens(input_text, tokenizer, max_tokens)
|
181 |
-
for chunk in chunks:
|
182 |
-
generated_text += model.generate(chunk)
|
183 |
-
return generated_text
|
184 |
-
|
185 |
-
@app.post("/generate")
|
186 |
-
async def generate_text(model_name: str = Query(...), input_text: str = Query(...)):
|
187 |
-
try:
|
188 |
-
model_loader = S3DirectStream(S3_BUCKET_NAME)
|
189 |
-
model = await model_loader.load_model_from_s3(model_name)
|
190 |
-
tokenizer = await model_loader.load_tokenizer_from_s3(model_name)
|
191 |
-
|
192 |
-
chunks = split_text_by_tokens(input_text, tokenizer, max_tokens=MAX_TOKENS)
|
193 |
-
|
194 |
-
generated_text = continue_generation(input_text, model, tokenizer)
|
195 |
-
|
196 |
-
return {"generated_text": generated_text}
|
197 |
-
|
198 |
-
except Exception as e:
|
199 |
-
raise HTTPException(status_code=500, detail=str(e))
|
200 |
-
|
201 |
-
if __name__ == "__main__":
|
202 |
-
import uvicorn
|
203 |
-
uvicorn.run(app, host="0.0.0.0", port=8000)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|