367 lines
12 KiB
Python
367 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import re
|
||
import urllib.request
|
||
import struct
|
||
from collections import Counter
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import yaml
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
MODEL_CARDS_DIR = ROOT / "registry" / "model-cards"
|
||
TOKEN_RE = re.compile(r"[A-Za-zА-Яа-яЁё0-9_]+", re.UNICODE)
|
||
RUSSIAN_ENDINGS = (
|
||
"иями",
|
||
"ями",
|
||
"ами",
|
||
"ого",
|
||
"ему",
|
||
"ыми",
|
||
"ими",
|
||
"ой",
|
||
"ей",
|
||
"ых",
|
||
"их",
|
||
"ую",
|
||
"юю",
|
||
"ая",
|
||
"яя",
|
||
"ое",
|
||
"ее",
|
||
"ом",
|
||
"ем",
|
||
"ам",
|
||
"ям",
|
||
"ах",
|
||
"ях",
|
||
"ы",
|
||
"и",
|
||
"а",
|
||
"я",
|
||
"е",
|
||
"у",
|
||
"ю",
|
||
)
|
||
TOKEN_ALIASES = {
|
||
"1с": ["1c", "bsl", "конфигурация"],
|
||
"1c": ["1с", "bsl", "configuration"],
|
||
"бсл": ["bsl", "1с"],
|
||
"bsl": ["бсл", "1с"],
|
||
"справочник": ["catalog", "справочники"],
|
||
"справочники": ["справочник", "catalog"],
|
||
"документ": ["documents", "документы"],
|
||
"документы": ["документ", "documents"],
|
||
"регистр": ["register", "регистры"],
|
||
"регистры": ["регистр", "register"],
|
||
"реквизит": ["attribute", "реквизиты"],
|
||
"реквизиты": ["реквизит", "attribute"],
|
||
"табличная": ["табличные", "часть"],
|
||
"табличные": ["табличная", "часть"],
|
||
"запрос": ["query", "read", "select", "выбрать"],
|
||
"выбрать": ["запрос", "query", "select"],
|
||
"форма": ["forms", "управляемая"],
|
||
"модуль": ["module", "bsl"],
|
||
"метаданные": ["metadata", "схема", "snapshot"],
|
||
"схема": ["metadata", "метаданные"],
|
||
"номенклатура": ["справочник", "catalog"],
|
||
}
|
||
|
||
|
||
def read_json(path: Path) -> dict:
|
||
with path.open("r", encoding="utf-8-sig") as handle:
|
||
data = json.load(handle)
|
||
if not isinstance(data, dict):
|
||
raise ValueError(f"{path} must contain a JSON object")
|
||
return data
|
||
|
||
|
||
def read_jsonl(path: Path) -> list[dict]:
|
||
if not path.exists():
|
||
return []
|
||
|
||
records: list[dict] = []
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
for line_number, line in enumerate(handle, start=1):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
try:
|
||
record = json.loads(line)
|
||
except json.JSONDecodeError as exc:
|
||
raise ValueError(f"{path}:{line_number}: invalid JSONL: {exc}") from exc
|
||
if not isinstance(record, dict):
|
||
raise ValueError(f"{path}:{line_number}: record must be an object")
|
||
records.append(record)
|
||
return records
|
||
|
||
|
||
def write_json(path: Path, data: dict) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||
|
||
|
||
def read_yaml_mapping(path: Path) -> dict[str, Any]:
|
||
with path.open("r", encoding="utf-8") as handle:
|
||
data = yaml.safe_load(handle)
|
||
if not isinstance(data, dict):
|
||
raise ValueError(f"{path} must contain a YAML mapping")
|
||
return data
|
||
|
||
|
||
def iter_model_card_paths(*, include_examples: bool = False) -> list[Path]:
|
||
paths = sorted(MODEL_CARDS_DIR.rglob("*.yaml")) + sorted(MODEL_CARDS_DIR.rglob("*.yml"))
|
||
if include_examples:
|
||
return paths
|
||
return [path for path in paths if "examples" not in path.parts]
|
||
|
||
|
||
def load_model_card(card_id: str) -> dict[str, Any]:
|
||
for suffix in (".yaml", ".yml"):
|
||
path = MODEL_CARDS_DIR / f"{card_id}{suffix}"
|
||
if path.exists():
|
||
return read_yaml_mapping(path)
|
||
raise FileNotFoundError(f"Model card not found for id `{card_id}` in {MODEL_CARDS_DIR}")
|
||
|
||
|
||
def localize_workspace_path(path: str) -> Path:
|
||
if path.startswith("/workspace/"):
|
||
if Path("/workspace").exists():
|
||
return Path(path)
|
||
return ROOT / path.removeprefix("/workspace/")
|
||
if path.startswith("/models/"):
|
||
if Path("/models").exists():
|
||
return Path(path)
|
||
return ROOT / "models" / path.removeprefix("/models/")
|
||
return Path(path)
|
||
|
||
|
||
def stem_russian_token(token: str) -> str:
|
||
if not re.search(r"[а-яё]", token, flags=re.IGNORECASE) or len(token) < 6:
|
||
return token
|
||
for ending in RUSSIAN_ENDINGS:
|
||
if token.endswith(ending) and len(token) - len(ending) >= 4:
|
||
return token[: -len(ending)]
|
||
return token
|
||
|
||
|
||
def normalize_token(token: str) -> str:
|
||
token = token.lower().replace("ё", "е")
|
||
return stem_russian_token(token)
|
||
|
||
|
||
def tokenize(text: str, *, expand_aliases: bool = False) -> list[str]:
|
||
tokens: list[str] = []
|
||
for raw_token in TOKEN_RE.findall(text):
|
||
token = normalize_token(raw_token)
|
||
tokens.append(token)
|
||
if expand_aliases:
|
||
tokens.extend(normalize_token(alias) for alias in TOKEN_ALIASES.get(token, []))
|
||
return tokens
|
||
|
||
|
||
def corpus_content_hash(records: list[dict]) -> str:
|
||
import hashlib
|
||
|
||
hasher = hashlib.sha256()
|
||
for record in records:
|
||
stable = {
|
||
"id": record.get("id"),
|
||
"document_id": record.get("document_id"),
|
||
"source_path": record.get("source_path"),
|
||
"chunk_index": record.get("chunk_index"),
|
||
"title": record.get("title"),
|
||
"content": record.get("content"),
|
||
"metadata": record.get("metadata") or {},
|
||
}
|
||
hasher.update(json.dumps(stable, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8"))
|
||
hasher.update(b"\n")
|
||
return hasher.hexdigest()
|
||
|
||
|
||
def hashing_embedding(text: str, *, dimensions: int = 384) -> list[float]:
|
||
import hashlib
|
||
|
||
if dimensions < 8:
|
||
raise ValueError("dimensions must be >= 8")
|
||
vector = [0.0] * dimensions
|
||
tokens = tokenize(text, expand_aliases=True)
|
||
if not tokens:
|
||
return vector
|
||
counts = Counter(tokens)
|
||
for token, count in counts.items():
|
||
digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
|
||
bucket = int.from_bytes(digest[:4], "little") % dimensions
|
||
sign = 1.0 if digest[4] & 1 else -1.0
|
||
vector[bucket] += sign * (1.0 + math.log(float(count)))
|
||
norm = math.sqrt(sum(value * value for value in vector))
|
||
if norm <= 0:
|
||
return vector
|
||
return [value / norm for value in vector]
|
||
|
||
|
||
def pack_float_vector(vector: list[float]) -> bytes:
|
||
return struct.pack(f"<{len(vector)}f", *vector)
|
||
|
||
|
||
def unpack_float_vector(data: bytes, dimensions: int) -> list[float]:
|
||
expected_size = dimensions * 4
|
||
if len(data) != expected_size:
|
||
raise ValueError(f"Vector blob has {len(data)} bytes, expected {expected_size}")
|
||
return list(struct.unpack(f"<{dimensions}f", data))
|
||
|
||
|
||
def cosine_similarity(left: list[float], right: list[float]) -> float:
|
||
if not left or not right or len(left) != len(right):
|
||
return 0.0
|
||
return float(sum(a * b for a, b in zip(left, right)))
|
||
|
||
|
||
def build_lexical_index(records: list[dict]) -> dict:
|
||
documents = []
|
||
document_frequency: Counter[str] = Counter()
|
||
total_length = 0
|
||
|
||
for record in records:
|
||
content = record.get("content") or ""
|
||
tokens = tokenize(content)
|
||
title_tokens = tokenize(str(record.get("title") or ""))
|
||
term_frequency = Counter(tokens)
|
||
document_frequency.update(term_frequency.keys())
|
||
total_length += len(tokens)
|
||
documents.append(
|
||
{
|
||
"id": record.get("id"),
|
||
"document_id": record.get("document_id"),
|
||
"source_path": record.get("source_path"),
|
||
"title": record.get("title"),
|
||
"chunk_index": record.get("chunk_index"),
|
||
"content": content,
|
||
"metadata": record.get("metadata") or {},
|
||
"length": len(tokens),
|
||
"title_tokens": title_tokens,
|
||
"term_frequency": dict(term_frequency),
|
||
}
|
||
)
|
||
|
||
doc_count = len(documents)
|
||
idf = {
|
||
token: math.log((1 + doc_count) / (1 + frequency)) + 1
|
||
for token, frequency in document_frequency.items()
|
||
}
|
||
|
||
return {
|
||
"schema_version": 2,
|
||
"type": "lexical-bm25",
|
||
"doc_count": doc_count,
|
||
"avg_doc_length": round(total_length / doc_count, 4) if doc_count else 0,
|
||
"idf": idf,
|
||
"documents": documents,
|
||
}
|
||
|
||
|
||
def score_lexical_document(query_tf: Counter[str], document: dict, idf: dict, avg_doc_length: float) -> float:
|
||
doc_tf = document.get("term_frequency") or {}
|
||
doc_length = max(float(document.get("length") or 0), 1.0)
|
||
avg_doc_length = max(float(avg_doc_length or doc_length), 1.0)
|
||
title_tokens = set(document.get("title_tokens") or [])
|
||
score = 0.0
|
||
k1 = 1.4
|
||
b = 0.72
|
||
for token, query_count in query_tf.items():
|
||
doc_count = doc_tf.get(token, 0)
|
||
if doc_count:
|
||
numerator = doc_count * (k1 + 1)
|
||
denominator = doc_count + k1 * (1 - b + b * (doc_length / avg_doc_length))
|
||
score += query_count * float(idf.get(token, 1.0)) * (numerator / denominator)
|
||
if token in title_tokens:
|
||
score += 0.35 * query_count
|
||
return score
|
||
|
||
|
||
def search_lexical_index(
|
||
index: dict,
|
||
query: str,
|
||
limit: int,
|
||
*,
|
||
candidate_limit: int | None = None,
|
||
dedupe_by_document: bool = False,
|
||
min_score: float = 0.0,
|
||
source_types: list[str] | None = None,
|
||
metadata_filters: dict[str, str] | None = None,
|
||
) -> list[dict]:
|
||
query_tf = Counter(tokenize(query, expand_aliases=True))
|
||
if not query_tf:
|
||
return []
|
||
|
||
results = []
|
||
idf = index.get("idf") or {}
|
||
avg_doc_length = float(index.get("avg_doc_length") or 0)
|
||
allowed_source_types = {source_type.lower() for source_type in source_types or []}
|
||
exact_filters = {key: str(value).lower() for key, value in (metadata_filters or {}).items() if str(value).strip()}
|
||
for document in index.get("documents") or []:
|
||
metadata = document.get("metadata") or {}
|
||
source_type = str(metadata.get("source_type") or "").lower()
|
||
if allowed_source_types and source_type not in allowed_source_types:
|
||
continue
|
||
if exact_filters and any(str(metadata.get(key) or "").lower() != value for key, value in exact_filters.items()):
|
||
continue
|
||
score = score_lexical_document(query_tf, document, idf, avg_doc_length)
|
||
if score > min_score:
|
||
results.append({"score": score, "document": document})
|
||
|
||
results.sort(key=lambda item: item["score"], reverse=True)
|
||
if candidate_limit:
|
||
results = results[:candidate_limit]
|
||
if dedupe_by_document:
|
||
deduped = []
|
||
seen = set()
|
||
for result in results:
|
||
document_id = result["document"].get("document_id") or result["document"].get("source_path")
|
||
if document_id in seen:
|
||
continue
|
||
seen.add(document_id)
|
||
deduped.append(result)
|
||
results = deduped
|
||
return results[:limit]
|
||
|
||
|
||
def call_chat_completion(
|
||
base_url: str,
|
||
model: str,
|
||
messages: list[dict],
|
||
*,
|
||
temperature: float = 0.2,
|
||
max_tokens: int = 1000,
|
||
timeout: int = 180,
|
||
) -> str:
|
||
url = f"{base_url.rstrip('/')}/v1/chat/completions"
|
||
payload = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens,
|
||
}
|
||
request = urllib.request.Request(
|
||
url,
|
||
data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||
data = json.loads(response.read().decode("utf-8"))
|
||
|
||
choices = data.get("choices") or []
|
||
if not choices:
|
||
raise ValueError("chat response has no choices")
|
||
|
||
message = choices[0].get("message") or {}
|
||
content = (message.get("content") or "").strip()
|
||
if not content:
|
||
raise ValueError("chat response content is empty")
|
||
return content
|