Update main.py

main.py

@@ -1,53 +1,104 @@
+# from fastapi import FastAPI
+# from pydantic import BaseModel
+# from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
+# import torch
+
+# app = FastAPI()
+
+# model_id = "HuggingFaceTB/SmolLM2-360M"
+# tokenizer = AutoTokenizer.from_pretrained(model_id)
+# model = AutoModelForCausalLM.from_pretrained(model_id)
+
+# class ChatRequest(BaseModel):
+#     context: str  # Historial de la conversación, como texto
+
+# class NewlineStoppingCriteria(StoppingCriteria):
+#     def __init__(self, prompt_len, tokenizer):
+#         super().__init__()
+#         self.prompt_len = prompt_len
+#         self.tokenizer = tokenizer
+
+#     def __call__(self, input_ids, scores, **kwargs):
+#         # Chequea si después del prompt hay un token de salto de línea
+#         gen_tokens = input_ids[0][self.prompt_len:]
+#         gen_text = self.tokenizer.decode(gen_tokens, skip_special_tokens=True)
+#         return '\n' in gen_text
+
+# @app.post("/chat/demo_base")
+# async def chat_demo_base(request: ChatRequest):
+#     prompt = (
+#         "Conversacion 1:\n"
+#         "-Dauro: -Hola Juanjo.\n"
+#         "-Juanjo: -¿Qué tal?\n"
+#         "-Dauro: -Bien, ¿y tú?\n\n"
+#         "Conversacion 2:\n"
+#         "-Juanjo: -Oye Asistente, ¿puedes mirar esto?\n"
+#         "-Asistente: -Por supuesto, dime.\n\n"
+#         f"Conversacion 3:\n{request.context}\n"
+#     )
+
+#     inputs = tokenizer(prompt, return_tensors="pt")
+#     input_ids = inputs["input_ids"]
+#     attention_mask = inputs["attention_mask"]
+
+#     stopping_criteria = StoppingCriteriaList([
+#         NewlineStoppingCriteria(prompt_len=input_ids.shape[1], tokenizer=tokenizer)
+#     ])
+
+#     output = model.generate(
+#         input_ids=input_ids,
+#         attention_mask=attention_mask,
+#         max_new_tokens=15,
+#         temperature=0.9,
+#         top_p=0.8,
+#         do_sample=True,
+#         pad_token_id=tokenizer.eos_token_id if hasattr(tokenizer, "eos_token_id") else None,
+#         stopping_criteria=stopping_criteria,
+#     )
+
+#     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+#     # Solo el fragmento después del prompt
+#     continuation = generated_text[len(prompt):].split('\n')[0]
+
+#     return {"generated_text": generated_text}
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from fastapi.responses import StreamingResponse
-import torch
-import threading
+from typing import List, Optional
 
 app = FastAPI()
 
-# Cargar modelo y tokenizer de Phi-2 (usa el modelo de Hugging Face Hub)
-model_id = "HuggingFaceTB/SmolLM2-135M"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+# Almacenamiento en memoria temporal
+registro_actual = {}
 
 # Modelo de entrada
-class ChatRequest(BaseModel):
-    message: str
-
-@app.post("/chat/stream")
-async def chat_stream(request: ChatRequest):
-    prompt = f"""Responde en español de forma clara y breve como un asistente IA.
-Usuario: {request.message}
-IA:"""
-
-    # Tokenizar entrada
-    inputs = tokenizer(prompt, return_tensors="pt")
-    input_ids = inputs["input_ids"]
-    attention_mask = inputs["attention_mask"]
-
-    # Streamer para obtener tokens generados poco a poco
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=False, skip_special_tokens=False)
-
-    # Iniciar la generación en un hilo aparte
-    generation_kwargs = dict(
-        input_ids=input_ids,    
-        attention_mask=attention_mask,
-        max_new_tokens=48,  # Puedes ajustar este valor para más/menos tokens
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        streamer=streamer,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    # StreamingResponse espera un generador que devuelva texto
-    async def event_generator():
-        for new_text in streamer:
-            yield new_text
-
-    return StreamingResponse(event_generator(), media_type="text/plain")
+class DialogoEntrada(BaseModel):
+    enunciado: str
+    personajes: List[str]  # lista de 3 personajes
+    relato_inicial: str
+    final_1: str
+    final_2: str
+    final_3: str
+
+# Modelo de salida
+class DialogoSalida(BaseModel):
+    enunciado: str
+    personajes: List[str]
+    relato_inicial: str
+    final_1: str
+    final_2: str
+    final_3: str
+
+@app.post("/entrada")
+async def registrar_dialogo(dialogo: DialogoEntrada):
+    global registro_actual
+    registro_actual = dialogo.dict()  # Sobrescribe el contenido anterior
+    return {"status": "registro guardado"}
 
+@app.get("/salida", response_model=Optional[DialogoSalida])
+async def obtener_y_limpiar():
+    global registro_actual
+    if not registro_actual:
+        return None
+    salida = registro_actual
+    registro_actual = {}  # Limpia después de devolver
+    return salida