DauroCamilo commited on
Commit
b55152c
verified
1 Parent(s): 5078b39

microsoft/phi-2

Browse files
Files changed (1) hide show
  1. main.py +48 -66
main.py CHANGED
@@ -1,71 +1,53 @@
1
- from fastapi import FastAPI, HTTPException
2
  from pydantic import BaseModel
3
- from typing import Dict, List
4
- import uuid
 
 
5
 
6
  app = FastAPI()
7
 
8
- # Estructura de datos en memoria
9
- salas: Dict[str, Dict] = {}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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