Spaces:
Sleeping
Sleeping
microsoft/phi-2
Browse files
main.py
CHANGED
@@ -1,71 +1,53 @@
|
|
1 |
-
from fastapi import FastAPI
|
2 |
from pydantic import BaseModel
|
3 |
-
from
|
4 |
-
import
|
|
|
|
|
5 |
|
6 |
app = FastAPI()
|
7 |
|
8 |
-
#
|
9 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
10 |
|
11 |
-
# Modelos de datos
|
12 |
-
class Mensaje(BaseModel):
|
13 |
-
usuario_id: str
|
14 |
-
texto: str
|
15 |
-
|
16 |
-
class Jugador(BaseModel):
|
17 |
-
nombre: str
|
18 |
-
|
19 |
-
@app.post("/sala")
|
20 |
-
def crear_sala():
|
21 |
-
while True:
|
22 |
-
nueva_id = str(uuid.uuid4())[:8]
|
23 |
-
if nueva_id not in salas:
|
24 |
-
salas[nueva_id] = {
|
25 |
-
"mensajes": [],
|
26 |
-
"jugadores": {}
|
27 |
-
}
|
28 |
-
return {"ok": True, "sala_id": nueva_id}
|
29 |
-
|
30 |
-
@app.post("/sala/{sala_id}/registrar")
|
31 |
-
def registrar_usuario(sala_id: str, jugador: Jugador):
|
32 |
-
if sala_id not in salas:
|
33 |
-
raise HTTPException(status_code=404, detail="Sala no encontrada")
|
34 |
-
|
35 |
-
if len(salas[sala_id]["jugadores"]) >= 4:
|
36 |
-
raise HTTPException(status_code=409, detail="Sala llena")
|
37 |
-
|
38 |
-
usuario_id = str(uuid.uuid4())
|
39 |
-
salas[sala_id]["jugadores"][usuario_id] = jugador.nombre
|
40 |
-
return {"usuario_id": usuario_id, "nombre": jugador.nombre}
|
41 |
-
|
42 |
-
@app.post("/sala/{sala_id}/mensaje")
|
43 |
-
def enviar_mensaje(sala_id: str, mensaje: Mensaje):
|
44 |
-
if sala_id not in salas:
|
45 |
-
raise HTTPException(status_code=404, detail="Sala no encontrada")
|
46 |
-
if mensaje.usuario_id not in salas[sala_id]["jugadores"]:
|
47 |
-
raise HTTPException(status_code=403, detail="Usuario no registrado en la sala")
|
48 |
-
|
49 |
-
nombre = salas[sala_id]["jugadores"][mensaje.usuario_id]
|
50 |
-
salas[sala_id]["mensajes"].append({
|
51 |
-
"usuario": nombre,
|
52 |
-
"texto": mensaje.texto
|
53 |
-
})
|
54 |
-
return {"ok": True}
|
55 |
-
|
56 |
-
@app.get("/sala/{sala_id}/mensajes")
|
57 |
-
def obtener_mensajes(sala_id: str):
|
58 |
-
if sala_id not in salas:
|
59 |
-
raise HTTPException(status_code=404, detail="Sala no encontrada")
|
60 |
-
return salas[sala_id]["mensajes"]
|
61 |
-
|
62 |
-
@app.get("/sala/{sala_id}/jugadores")
|
63 |
-
def obtener_jugadores(sala_id: str):
|
64 |
-
if sala_id not in salas:
|
65 |
-
raise HTTPException(status_code=404, detail="Sala no encontrada")
|
66 |
-
|
67 |
-
jugadores = [
|
68 |
-
{"posicion": idx, "nombre": nombre}
|
69 |
-
for idx, (_, nombre) in enumerate(salas[sala_id]["jugadores"].items())
|
70 |
-
]
|
71 |
-
return jugadores
|
|
|
1 |
+
from fastapi import FastAPI
|
2 |
from pydantic import BaseModel
|
3 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
|
4 |
+
from fastapi.responses import StreamingResponse
|
5 |
+
import torch
|
6 |
+
import threading
|
7 |
|
8 |
app = FastAPI()
|
9 |
|
10 |
+
# Cargar modelo y tokenizer de Phi-2 (usa el modelo de Hugging Face Hub)
|
11 |
+
model_id = "microsoft/phi-2"
|
12 |
+
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
13 |
+
model = AutoModelForCausalLM.from_pretrained(model_id)
|
14 |
+
|
15 |
+
# Modelo de entrada
|
16 |
+
class ChatRequest(BaseModel):
|
17 |
+
message: str
|
18 |
+
|
19 |
+
@app.post("/chat/stream")
|
20 |
+
async def chat_stream(request: ChatRequest):
|
21 |
+
prompt = f"""Responde en espa帽ol de forma clara y breve como un asistente IA.
|
22 |
+
Usuario: {request.message}
|
23 |
+
IA:"""
|
24 |
+
|
25 |
+
# Tokenizar entrada
|
26 |
+
inputs = tokenizer(prompt, return_tensors="pt")
|
27 |
+
input_ids = inputs["input_ids"]
|
28 |
+
attention_mask = inputs["attention_mask"]
|
29 |
+
|
30 |
+
# Streamer para obtener tokens generados poco a poco
|
31 |
+
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
32 |
+
|
33 |
+
# Iniciar la generaci贸n en un hilo aparte
|
34 |
+
generation_kwargs = dict(
|
35 |
+
input_ids=input_ids,
|
36 |
+
attention_mask=attention_mask,
|
37 |
+
max_new_tokens=48, # Puedes ajustar este valor para m谩s/menos tokens
|
38 |
+
temperature=0.7,
|
39 |
+
top_p=0.9,
|
40 |
+
do_sample=True,
|
41 |
+
streamer=streamer,
|
42 |
+
pad_token_id=tokenizer.eos_token_id
|
43 |
+
)
|
44 |
+
thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
|
45 |
+
thread.start()
|
46 |
+
|
47 |
+
# StreamingResponse espera un generador que devuelva texto
|
48 |
+
async def event_generator():
|
49 |
+
for new_text in streamer:
|
50 |
+
yield new_text
|
51 |
+
|
52 |
+
return StreamingResponse(event_generator(), media_type="text/plain")
|
53 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|