129 lines
3.8 KiB
Python
129 lines
3.8 KiB
Python
# utils/rag_engine.py
|
|
|
|
import os
|
|
import json
|
|
import hashlib
|
|
import numpy as np
|
|
|
|
_EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
|
|
|
|
def _plugin_root():
|
|
return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
def rag_doc_dir():
|
|
return os.path.join(_plugin_root(), "rag_doc")
|
|
|
|
def _index_dir():
|
|
d = os.path.join(_plugin_root(), ".rag_index")
|
|
os.makedirs(d, exist_ok=True)
|
|
return d
|
|
|
|
def list_rag_txt_files():
|
|
root = rag_doc_dir()
|
|
if not os.path.isdir(root):
|
|
return []
|
|
return [f for f in os.listdir(root) if f.lower().endswith(".txt")]
|
|
|
|
def _read_text(fp):
|
|
with open(fp, "r", encoding="utf-8", errors="ignore") as f:
|
|
return f.read()
|
|
|
|
def _chunk_text(text, chunk_size=800, overlap=100):
|
|
# 단순 문자 길이 기반 청킹
|
|
chunks = []
|
|
i = 0
|
|
n = len(text)
|
|
while i < n:
|
|
chunk = text[i : i + chunk_size]
|
|
chunks.append(chunk)
|
|
i += (chunk_size - overlap)
|
|
return chunks
|
|
|
|
def _hash_for_file(path):
|
|
h = hashlib.sha256()
|
|
with open(path, "rb") as f:
|
|
while True:
|
|
b = f.read(1024 * 1024)
|
|
if not b:
|
|
break
|
|
h.update(b)
|
|
return h.hexdigest()
|
|
|
|
def _embed_sentences(sentences):
|
|
"""
|
|
1순위: sentence-transformers
|
|
2순위: scikit-learn TfidfVectorizer
|
|
3순위: 간단한 백오프(문자 unigram 빈도)
|
|
"""
|
|
# 1) sentence-transformers
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
model = SentenceTransformer(_EMBED_MODEL_NAME, device="cpu") # CPU에도 충분히 작음
|
|
emb = model.encode(sentences, show_progress_bar=False, convert_to_numpy=True, normalize_embeddings=True)
|
|
return emb
|
|
except Exception:
|
|
pass
|
|
|
|
# 2) scikit-learn TF-IDF
|
|
try:
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
vect = TfidfVectorizer(max_features=4096)
|
|
X = vect.fit_transform(sentences).astype(np.float32)
|
|
# L2 normalize
|
|
row_norm = np.sqrt((X.power(2)).sum(axis=1)).A1 + 1e-12
|
|
X = X.multiply(1.0 / row_norm[:, None]).A
|
|
return X
|
|
except Exception:
|
|
pass
|
|
|
|
# 3) 아주 간단한 백오프: a-z 빈도(26차원)
|
|
def char_feat(s):
|
|
s = s.lower()
|
|
vec = np.zeros(26, dtype=np.float32)
|
|
for ch in s:
|
|
if "a" <= ch <= "z":
|
|
vec[ord(ch) - 97] += 1.0
|
|
if np.linalg.norm(vec) > 0:
|
|
vec /= np.linalg.norm(vec)
|
|
return vec
|
|
return np.stack([char_feat(s) for s in sentences], axis=0)
|
|
|
|
|
|
def build_or_load_index(txt_filename):
|
|
"""
|
|
반환: dict {
|
|
"chunks": [str, ...],
|
|
"emb": np.ndarray (n_chunks, d),
|
|
"txt_path": fullpath
|
|
}
|
|
"""
|
|
src = os.path.join(rag_doc_dir(), txt_filename)
|
|
if not os.path.exists(src):
|
|
raise FileNotFoundError(f"RAG txt not found: {src}")
|
|
|
|
sig = _hash_for_file(src)
|
|
cache_fp = os.path.join(_index_dir(), f"{os.path.basename(src)}.{sig}.npz")
|
|
|
|
if os.path.exists(cache_fp):
|
|
data = np.load(cache_fp, allow_pickle=True)
|
|
return {"chunks": data["chunks"].tolist(), "emb": data["emb"].astype(np.float32), "txt_path": src}
|
|
|
|
text = _read_text(src)
|
|
chunks = _chunk_text(text, chunk_size=800, overlap=100)
|
|
emb = _embed_sentences(chunks).astype(np.float32)
|
|
|
|
np.savez_compressed(cache_fp, chunks=np.array(chunks, dtype=object), emb=emb)
|
|
return {"chunks": chunks, "emb": emb, "txt_path": src}
|
|
|
|
def _cosine_topk(mat, q, k=4):
|
|
# mat: (n, d), q: (d,)
|
|
s = mat @ q # 가정: 이미 L2 정규화
|
|
idx = np.argsort(-s)[:k]
|
|
return idx, s[idx]
|
|
|
|
def query_with_rag(index, query, k=4):
|
|
emb_q = _embed_sentences([query])[0].astype(np.float32)
|
|
idx, _ = _cosine_topk(index["emb"], emb_q, k=k)
|
|
picks = [index["chunks"][i] for i in idx]
|
|
return "\n---\n".join(picks)
|