193 lines
7 KiB
Python
193 lines
7 KiB
Python
"""Buscar en el saber de COFRE lo que hace falta para contestar una pregunta.
|
|
|
|
Esto es la mitad "recuperar" de RAG: dada una pregunta, devuelve los fragmentos
|
|
de los apuntes que mas se le parecen, para pasarselos al cerebro. El cerebro
|
|
construye la respuesta desde ESE texto, no desde lo que recuerde, que es lo que
|
|
evita que se invente flags.
|
|
|
|
## Por que se carga una vez y se guarda
|
|
|
|
El indice son 2.866 vectores. Cargarlo y encodear la pregunta cuesta, pero solo
|
|
la primera vez: el modelo y los vectores se quedan en memoria. Una consulta
|
|
despues es un producto escalar contra 2.866 filas, milisegundos.
|
|
|
|
## Por que no una base de datos vectorial
|
|
|
|
FAISS, Chroma y compañia son para millones de vectores. Con 2.866, numpy y un
|
|
producto escalar es mas rapido de arrancar, no añade una dependencia pesada, y
|
|
el indice es un jsonl que se puede abrir y leer con los ojos.
|
|
"""
|
|
import json
|
|
import os
|
|
|
|
AQUI = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
INDICE = os.path.join(AQUI, "datos", "saber.jsonl")
|
|
|
|
_modelo = None
|
|
_entradas = None
|
|
_matriz = None
|
|
|
|
|
|
def _carga():
|
|
global _modelo, _entradas, _matriz
|
|
if _entradas is not None:
|
|
return _entradas is not False
|
|
|
|
if not os.path.exists(INDICE):
|
|
_entradas = False
|
|
return False
|
|
|
|
import numpy as np
|
|
from model2vec import StaticModel
|
|
|
|
entradas, vectores = [], []
|
|
with open(INDICE, encoding="utf-8") as f:
|
|
for l in f:
|
|
try:
|
|
d = json.loads(l)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
v = d.pop("v", None)
|
|
if v is None:
|
|
continue
|
|
entradas.append(d)
|
|
vectores.append(v)
|
|
|
|
_entradas = entradas
|
|
_matriz = np.array(vectores, dtype="float32")
|
|
# normalizar una vez: asi la similitud es un simple producto escalar
|
|
normas = np.linalg.norm(_matriz, axis=1, keepdims=True)
|
|
_matriz = _matriz / np.clip(normas, 1e-9, None)
|
|
_modelo = StaticModel.from_pretrained("minishlab/potion-multilingual-128M")
|
|
return True
|
|
|
|
|
|
def _norm(t: str) -> str:
|
|
import re
|
|
import unicodedata
|
|
t = unicodedata.normalize("NFD", (t or "").lower())
|
|
t = "".join(c for c in t if unicodedata.category(c) != "Mn")
|
|
return re.sub(r"[^a-z0-9 ]", " ", t)
|
|
|
|
|
|
# El texto normalizado de cada entrada y su nombre, cacheados para el bono.
|
|
_texto_norm = None
|
|
_nombre_norm = None
|
|
|
|
|
|
def _bono_palabras(palabras):
|
|
"""Un vector de bonos, uno por entrada, segun cuantas palabras compartan."""
|
|
import numpy as np
|
|
global _texto_norm, _nombre_norm
|
|
if _texto_norm is None:
|
|
_texto_norm = [set(_norm(e["texto"]).split()) for e in _entradas]
|
|
_nombre_norm = [_norm(e.get("herramienta", "")) for e in _entradas]
|
|
bono = np.zeros(len(_entradas), dtype="float32")
|
|
for i, (palabras_e, nombre) in enumerate(zip(_texto_norm, _nombre_norm)):
|
|
comunes = len(palabras & palabras_e)
|
|
b = 0.08 * comunes # cada palabra compartida suma
|
|
if nombre and nombre in palabras: # y nombrar la herramienta, mucho
|
|
b += 0.5
|
|
bono[i] = b
|
|
return bono
|
|
|
|
|
|
def busca(pregunta: str, cuantos: int = 5, perfil: str = None) -> list:
|
|
"""Los fragmentos mas parecidos a la pregunta, mejor primero.
|
|
|
|
Cada uno lleva su similitud (0-1), de que herramienta y fichero viene, y el
|
|
texto. perfil filtra por carpeta de COFRE si se quiere acotar.
|
|
"""
|
|
if not _carga() or not pregunta.strip():
|
|
return []
|
|
import numpy as np
|
|
|
|
q = _modelo.encode([pregunta])[0]
|
|
q = q / max(float(np.linalg.norm(q)), 1e-9)
|
|
sim = _matriz @ q
|
|
|
|
# Impulso por palabras exactas (busqueda hibrida).
|
|
#
|
|
# Los embeddings estaticos fallan cuando la pregunta va en castellano y el
|
|
# texto en ingles —"cambiar permisos" no se parece a "change file mode
|
|
# bits"—. Pero si la pregunta NOMBRA la herramienta ("como uso chmod") o
|
|
# comparte palabras con el texto, eso es una señal fuerte que la similitud
|
|
# semantica sola no aprovecha. Se suma un bono por cada palabra de la
|
|
# pregunta que aparezca, y uno grande si coincide el nombre de la
|
|
# herramienta: asi "usa nmap para..." lleva nmap al primer puesto.
|
|
palabras = {w for w in _norm(pregunta).split() if len(w) > 3}
|
|
if palabras:
|
|
bono = _bono_palabras(palabras)
|
|
sim = sim + bono
|
|
|
|
orden = np.argsort(-sim)[: cuantos * 6]
|
|
salida, vistos = [], set()
|
|
for i in orden:
|
|
e = _entradas[i]
|
|
if perfil and e.get("perfil") != perfil:
|
|
continue
|
|
# sin duplicados: OSCP_APUNTES es copia de los apuntes de primer nivel.
|
|
# Se comparan los primeros 120 caracteres, que basta para reconocerlos.
|
|
firma = e["texto"][:120]
|
|
if firma in vistos:
|
|
continue
|
|
vistos.add(firma)
|
|
salida.append({**e, "sim": round(float(sim[i]), 3)})
|
|
if len(salida) >= cuantos:
|
|
break
|
|
return salida
|
|
|
|
|
|
def contexto(pregunta: str, cuantos: int = 5, minimo: float = 0.35) -> str:
|
|
"""Lo mismo, pero ya montado como texto para meterle al cerebro.
|
|
|
|
Se corta por debajo de `minimo` de similitud: pasarle al modelo fragmentos
|
|
que no vienen a cuento es lo que le hace mezclar herramientas que no tienen
|
|
nada que ver. Mejor darle poco y bueno que mucho y ruidoso.
|
|
"""
|
|
trozos = [t for t in busca(pregunta, cuantos) if t["sim"] >= minimo]
|
|
if not trozos:
|
|
return ""
|
|
lineas = ["De los apuntes de COFRE del usuario:\n"]
|
|
for t in trozos:
|
|
fuente = t.get("fichero") or f"{t['perfil']}/{t['herramienta']}"
|
|
lineas.append(f"--- {fuente} ---\n{t['texto']}\n")
|
|
return "\n".join(lineas)
|
|
|
|
|
|
def resumen(tope_temas: int = 10) -> str:
|
|
"""Un resumen de QUE hay indexado, para que el modelo sepa de que sabe.
|
|
|
|
Se deriva del indice de VERDAD, no de una lista fija en el prompt. Asi es
|
|
honesto en cualquier maquina: en la del usuario dice sus areas reales; en un
|
|
clon sin apuntes propios dira lo que haya (comandos de Linux) o nada. Lo usa
|
|
el cerebro para el prompt de sistema; evita que presuma de un conocimiento
|
|
que en esa maquina no tiene."""
|
|
if not _carga() or not _entradas:
|
|
return ""
|
|
from collections import Counter
|
|
perfiles, temas = Counter(), Counter()
|
|
n_comandos = 0
|
|
for e in _entradas:
|
|
t = e.get("tipo")
|
|
if t == "comando":
|
|
n_comandos += 1
|
|
elif t in ("apunte", "metodologia", "ficha", "pregunta") and e.get("perfil"):
|
|
perfiles[e["perfil"]] += 1
|
|
if e.get("tema"):
|
|
temas[e["tema"]] += 1
|
|
partes = []
|
|
if perfiles:
|
|
partes.append("apuntes de " + ", ".join(p for p, _ in perfiles.most_common(6)))
|
|
if temas:
|
|
partes.append("metodologia en " + ", ".join(t for t, _ in temas.most_common(tope_temas)))
|
|
if n_comandos:
|
|
partes.append(f"{n_comandos} comandos de Linux")
|
|
if not partes:
|
|
return ""
|
|
return ("En los apuntes que tienes indexados (buscalos con buscar_en_apuntes) "
|
|
"hay: " + "; ".join(partes) + ".")
|
|
|
|
|
|
def disponible() -> bool:
|
|
return os.path.exists(INDICE)
|