atareao / rag_embeddings.py

"#!/usr/bin/env python3\n\"\"\"\nrag_embeddings.py — Generación de embeddings con Ollama + bge-m3.\n\nEnvía texto al modelo de embeddings de Ollama y empaqueta los vectores\nresultantes como BLOBs binarios para almacenar en SQLite.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport struct\nfrom typing import Any\n\nimport ollama\n\nfrom rag_config import MODELO_EMBEDDINGS, DIMENSIONES_EMBEDDING\n\n\n# ── Funciones de embedding ───────────────────────────────────────────────────\n\n\ndef get_embedding(texto: str, modelo: str | None = None) -> list[float]:\n \"\"\"Genera el embedding de un texto usando el modelo de Ollama.\n\n Args:\n texto: Texto a convertir en vector.\n modelo: Nombre del modelo en Ollama (por defecto bge-m3).\n\n Returns:\n Lista de floats con el vector de embedding (1024 dimensiones).\n \"\"\"\n resp: Any = ollama.embed(model=modelo or MODELO_EMBEDDINGS, input=texto)\n # La API devuelve una lista de embeddings; para un solo input\n # el primer (y único) elemento es el que nos interesa.\n return resp[\"embeddings\"][0]\n\n\ndef get_embeddings_batch(\n textos: list[str],\n batch_size: int = 16,\n modelo: str | None = None,\n) -> list[list[float]]:\n \"\"\"Genera embeddings en lotes, que es más eficiente que llamadas individuales.\n\n Args:\n textos: Lista de textos a embedder.\n batch_size: Tamaño del lote para cada llamada a Ollama.\n modelo: Nombre del modelo en Ollama.\n\n Returns:\n Lista de embeddings (cada uno es una lista de 1024 floats).\n \"\"\"\n resultados: list[list[float]] = []\n modelo = modelo or MODELO_EMBEDDINGS\n\n for i in range(0, len(textos), batch_size):\n batch = textos[i : i + batch_size]\n resp: Any = ollama.embed(model=modelo, input=batch)\n resultados.extend(resp[\"embeddings\"])\n\n return resultados\n\n\n# ── Empaquetado binario ──────────────────────────────────────────────────────\n\n\ndef pack_vector(vec: list[float]) -> bytes:\n \"\"\"Convierte una lista de floats a un BLOB binario (4 bytes por float).\n\n Args:\n vec: Lista de floats (1024 dimensiones para bge-m3).\n\n Returns:\n Bytes empaquetados con struct.pack (little-endian, float32).\n \"\"\"\n return struct.pack(f\"{len(vec)}f\", *vec)\n\n\ndef unpack_vector(data: bytes) -> list[float]:\n \"\"\"Convierte un BLOB binario de vuelta a una lista de floats.\n\n Args:\n data: Bytes empaquetados (de pack_vector).\n\n Returns:\n Lista de floats.\n \"\"\"\n # El formato 'f' usa float32 (4 bytes)\n n = len(data) // 4\n return list(struct.unpack(f\"{n}f\", data))\n\n\n# ── CLI de prueba ────────────────────────────────────────────────────────────\n\n\ndef main() -> None:\n \"\"\"Prueba rápida: genera un embedding y muestra estadísticas.\"\"\"\n texto = \"¿Cómo configuro systemd timers en Linux?\"\n print(f'Generando embedding para: \"{texto}\"\\n')\n\n embedding = get_embedding(texto)\n print(f\"Dimensiones: {len(embedding)}\")\n print(f\"Primeros 5 valores: {embedding[:5]}\")\n print(f\"Norma L2: {sum(x * x for x in embedding) ** 0.5:.6f}\")\n\n # Probar empaquetado\n blob = pack_vector(embedding)\n print(f\"\\nTamaño del BLOB: {len(blob)} bytes (esperado: {len(embedding) * 4})\")\n\n recovered = unpack_vector(blob)\n print(f\"Recuperados: {len(recovered)} floats\")\n print(f\"Coinciden: {abs(embedding[0] - recovered[0]) < 1e-6}\")\n\n\nif __name__ == \"__main__\":\n main()\n"