mabosaimi/text2sql

1from __future__ import annotations2 3import json4from pathlib import Path5from typing import List, Tuple, Dict, Any6 7import torch8from huggingface_hub import hf_hub_download9from sentence_transformers import SentenceTransformer10from sentence_transformers.util import cos_sim11 12_MODEL_ID = "mabosaimi/bge-m3-text2tables"13 14model: SentenceTransformer = SentenceTransformer(_MODEL_ID)15 16_corpus_text_file = hf_hub_download(repo_id=_MODEL_ID, filename="corpus_texts.json")17with open(_corpus_text_file, "r", encoding="utf-8") as _f:18    corpus_texts: List[str] = json.load(_f)19 20_corpus_emb_file = hf_hub_download(repo_id=_MODEL_ID, filename="corpus_embeddings.pt")21corpus_embeddings: torch.Tensor = torch.load(_corpus_emb_file, map_location="cpu")22 23_schemas_PATH = Path(__file__).parent / "schemas.json"24if _schemas_PATH.exists():25    with open(_schemas_PATH, "r", encoding="utf-8") as _cf:26        schemas: List[Dict[str, Any]] = json.load(_cf)27else:28    schemas = []29 30 31def get_model_id() -> str:32    """Return the identifier of the embedding model in use.33 34    This intentionally hides low-level model details from API consumers while35    allowing health/diagnostics endpoints to expose basic service info.36    """37 38    return _MODEL_ID39 40 41def get_corpus_size() -> int:42    """Return the number of entries in the fixed metadata corpus."""43 44    return len(corpus_texts)45 46 47def preprocess_text(query: str) -> str:48    """Preprocess a natural language string by stripping whitespace.49 50    Inputs:51    - query: Natural language string to be preprocessed.52 53    Returns:54    - The preprocessed string.55    """56    return query.strip()57 58 59def encode_text(query: str) -> torch.Tensor:60    """Encode a natural language query into an embedding tensor.61 62    Inputs:63    - query: Natural language string to be embedded.64 65    Returns:66    - A 1 x D torch.Tensor representing the normalized embedding of the query.67    """68    query = preprocess_text(query)69    return model.encode(query, convert_to_tensor=True, normalize_embeddings=True)70 71 72def semantic_search(query: str, top_k: int = 5) -> List[Tuple[float, str, int]]:73    """Compute semantic similarity between a query and the stored corpus.74 75    Inputs:76    - query: Natural language search string.77    - top_k: Maximum number of results to return (capped at corpus size).78 79    Returns:80    - A list of tuples (score, text, index) sorted by descending similarity,81      where:82        - score is a float cosine similarity.83        - text is the matched corpus entry.84        - index is the integer position in the corpus (stable identifier).85    """86 87    query_embedding = encode_text(query)88    scores = cos_sim(query_embedding, corpus_embeddings)[0]89    k = min(max(top_k, 1), len(corpus_texts))90    values, indices = torch.topk(scores, k=k)91    return [92        (float(values[i]), corpus_texts[int(indices[i])], int(indices[i]))93        for i in range(len(values))94    ]95 96 97def get_schemas(include_columns: bool = False) -> List[Dict[str, Any]]:98    """Return the local schemas.99 100    Inputs:101    - include_columns: When True, include full column metadata; otherwise102      return a minimal view with table name and description only.103 104    Returns:105    - List of table dicts. If include_columns is False, each dict contains106      {"table", "description"}. If True, it includes the original structure.107    """108 109    if not schemas:110        return []111    if include_columns:112        return schemas113    return [114        {"table": t["table"], "description": t.get("description", "")} for t in schemas115    ]116