app.py

import os
import gc
import torch
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, JSONResponse
from pydantic import BaseModel
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig
)
import uvicorn
from duckduckgo_search import ddg  # pip install duckduckgo-search

# Nombre del modelo fijo
MODEL_NAME = "lilmeaty/my_xdd"

# Variables globales para almacenar el modelo y el tokenizador
global_model = None
global_tokenizer = None

# Función asíncrona para limpiar la memoria (RAM y caché CUDA)
async def cleanup_memory(device: str):
    gc.collect()
    if device == "cuda":
        torch.cuda.empty_cache()
    await asyncio.sleep(0.01)

# Request para la generación de texto
class GenerateRequest(BaseModel):
    input_text: str = ""
    max_new_tokens: int = 200  # límite total de tokens generados (puede ser muy alto)
    temperature: float = 1.0
    top_p: float = 1.0
    top_k: int = 50
    repetition_penalty: float = 1.0
    do_sample: bool = True
    stream: bool = False
    # Límite de tokens por chunk en modo streaming (si se excede, se emite el chunk acumulado)
    chunk_token_limit: int = 20
    # Secuencias que, si se detectan, hacen que se detenga la generación
    stop_sequences: list[str] = []
    # Si se desea incluir Duckasgo en la respuesta final
    include_duckasgo: bool = False

# Request para búsquedas independientes con Duckasgo
class DuckasgoRequest(BaseModel):
    query: str

# Inicializar la aplicación FastAPI
app = FastAPI()

# Evento de startup: cargar el modelo y tokenizador globalmente
@app.on_event("startup")
async def load_global_model():
    global global_model, global_tokenizer
    try:
        config = AutoConfig.from_pretrained(MODEL_NAME)
        global_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, config=config)
        # Se usa torch_dtype=torch.float16 para reducir la huella de memoria (si es posible)
        global_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, config=config, torch_dtype=torch.float16)
        # Configurar token de padding si es necesario
        if global_tokenizer.eos_token_id is not None and global_tokenizer.pad_token_id is None:
            global_tokenizer.pad_token_id = config.pad_token_id or global_tokenizer.eos_token_id
        device = "cuda" if torch.cuda.is_available() else "cpu"
        global_model.to(device)
        print(f"Modelo {MODEL_NAME} cargado correctamente en {device}.")
    except Exception as e:
        print("Error al cargar el modelo:", e)

# Función para realizar búsqueda con Duckasgo de forma asíncrona
async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
    results = await asyncio.to_thread(ddg, query, max_results=max_results)
    if not results:
        return "No se encontraron resultados en Duckasgo."
    summary = "\nResultados de búsqueda (Duckasgo):\n"
    for idx, res in enumerate(results, start=1):
        title = res.get("title", "Sin título")
        url = res.get("href", "Sin URL")
        snippet = res.get("body", "")
        summary += f"{idx}. {title}\n   URL: {url}\n   {snippet}\n"
    return summary

# Función para generar texto en modo streaming, dividiendo en chunks ilimitados
# y deteniendo la generación si se detectan secuencias de parada.
async def stream_text(request: GenerateRequest, device: str):
    global global_model, global_tokenizer

    # Codificar la entrada y obtener la longitud inicial
    encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
    initial_input_len = encoded_input.input_ids.shape[-1]
    input_ids = encoded_input.input_ids

    # Variables para acumular texto
    accumulated_text = ""  # todo el texto generado
    current_chunk = ""     # chunk actual
    chunk_token_count = 0

    # Configurar la generación según los parámetros recibidos
    gen_config = GenerationConfig(
        temperature=request.temperature,
        max_new_tokens=request.max_new_tokens,
        top_p=request.top_p,
        top_k=request.top_k,
        repetition_penalty=request.repetition_penalty,
        do_sample=request.do_sample,
    )

    past_key_values = None

    # Usamos un bucle while para permitir generación ilimitada en chunks
    for _ in range(request.max_new_tokens):
        with torch.no_grad():
            outputs = global_model(
                input_ids,
                past_key_values=past_key_values,
                use_cache=True,
                return_dict=True
            )
        logits = outputs.logits[:, -1, :]
        past_key_values = outputs.past_key_values

        if gen_config.do_sample:
            logits = logits / gen_config.temperature
            if gen_config.top_k and gen_config.top_k > 0:
                topk_values, _ = torch.topk(logits, k=gen_config.top_k)
                logits[logits < topk_values[:, [-1]]] = -float('Inf')
            probs = torch.nn.functional.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
        else:
            next_token = torch.argmax(logits, dim=-1, keepdim=True)

        token_id = next_token.item()
        token_text = global_tokenizer.decode([token_id], skip_special_tokens=True)

        accumulated_text += token_text
        current_chunk += token_text
        chunk_token_count += 1

        # Verificar si se alcanzó alguna de las secuencias de parada
        if request.stop_sequences:
            for stop_seq in request.stop_sequences:
                if stop_seq in accumulated_text:
                    # Si se detecta el stop, emitir el chunk actual y terminar
                    yield current_chunk
                    await cleanup_memory(device)
                    return

        # Si se supera el límite de tokens por chunk, enviar el chunk acumulado
        if chunk_token_count >= request.chunk_token_limit:
            yield current_chunk
            current_chunk = ""
            chunk_token_count = 0

        # Permitir que otras tareas se ejecuten
        await asyncio.sleep(0)

        # Actualizar el input para la siguiente iteración
        input_ids = next_token

        # Si se ha generado el token de finalización, se emite el chunk y se termina
        if token_id == global_tokenizer.eos_token_id:
            break

    # Emitir el último chunk (si no está vacío)
    if current_chunk:
        yield current_chunk

    # Si se solicitó incluir Duckasgo, realizar la búsqueda y agregarla como chunk final
    if request.include_duckasgo:
        search_summary = await perform_duckasgo_search(request.input_text)
        yield "\n" + search_summary

    await cleanup_memory(device)

# Endpoint para la generación de texto (modo streaming o no-streaming)
@app.post("/generate")
async def generate_text(request: GenerateRequest):
    global global_model, global_tokenizer
    if global_model is None or global_tokenizer is None:
        raise HTTPException(status_code=500, detail="El modelo no se ha cargado correctamente.")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    gen_config = GenerationConfig(
        temperature=request.temperature,
        max_new_tokens=request.max_new_tokens,
        top_p=request.top_p,
        top_k=request.top_k,
        repetition_penalty=request.repetition_penalty,
        do_sample=request.do_sample,
    )

    try:
        if request.stream:
            # Modo streaming: se envían múltiples chunks según se generan
            generator = stream_text(request, device)
            return StreamingResponse(generator, media_type="text/plain")
        else:
            # Modo no-streaming: generación completa en una sola respuesta
            encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
            with torch.no_grad():
                output = global_model.generate(
                    **encoded_input,
                    generation_config=gen_config,
                    return_dict_in_generate=True,
                    output_scores=True
                )
            input_length = encoded_input["input_ids"].shape[-1]
            generated_text = global_tokenizer.decode(
                output.sequences[0][input_length:],
                skip_special_tokens=True
            )
            # Si se han definido secuencias de parada, se corta la respuesta en el primer match
            if request.stop_sequences:
                for stop_seq in request.stop_sequences:
                    if stop_seq in generated_text:
                        generated_text = generated_text.split(stop_seq)[0]
                        break
            # Si se solicitó incluir Duckasgo, se realiza la búsqueda y se agrega al final
            if request.include_duckasgo:
                search_summary = await perform_duckasgo_search(request.input_text)
                generated_text += "\n" + search_summary
            await cleanup_memory(device)
            return {"generated_text": generated_text}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error durante la generación: {e}")

# Endpoint independiente para búsquedas con Duckasgo
@app.post("/duckasgo")
async def duckasgo_search(request: DuckasgoRequest):
    try:
        results = await asyncio.to_thread(ddg, request.query, max_results=10)
        if results is None:
            results = []
        return JSONResponse(content={"query": request.query, "results": results})
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error en la búsqueda: {e}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)