src/metadata_llm.py

"""Auto-metadatos vía LLM HuggingFace."""
from transformers import pipeline
import os

# Por defecto usa Gemma 2B instruct, puedes cambiarlo.
LLM_MODEL = os.getenv("LLM_METADATA_MODEL", "google/gemma-1.1-2b-it")

# Sólo cargar el pipeline una vez (lazy).
_llm = None

def get_llm():
    global _llm
    if _llm is None:
        _llm = pipeline(
            "text-generation",
            model=LLM_MODEL,
            device_map="auto",
            max_new_tokens=256,
            trust_remote_code=True
        )
    return _llm

def build_prompt(text: str) -> str:
    # Puedes personalizar el prompt aquí
    return (
        "Analiza el siguiente texto y responde en JSON con las claves:\n"
        "\"summary\": resumen en una frase;\n"
        "\"topics\": lista de hasta 5 palabras clave o temas principales;\n"
        "\"language\": idioma detectado ('es', 'en', etc.);\n"
        "\"vertical\": vertical temática (SEO-LLM, eCommerce, etc).\n\n"
        "TEXTO:\n" + text.strip() + "\n\nJSON:"
    )

def extract_metadata(text: str) -> dict:
    llm = get_llm()
    prompt = build_prompt(text)
    output = llm(prompt)[0]['generated_text']
    # Extraer sólo el bloque JSON del resultado
    import re, json
    m = re.search(r'\\{.*\\}', output, re.DOTALL)
    if not m:
        return {}
    try:
        result = json.loads(m.group(0))
        return result
    except Exception:
        return {}

# --- CLI de prueba rápida
if __name__ == "__main__":
    chunk = "El SEO para LLMs consiste en optimizar contenido pensando en cómo los grandes modelos de lenguaje recuperan información relevante, mejorando la estructura, claridad y contexto del texto para aumentar la calidad de las respuestas generadas por IA."
    meta = extract_metadata(chunk)
    print(meta)