maikheb/nl2sql

1import faiss
2import os
3import pickle
4from sentence_transformers import SentenceTransformer
5import numpy as np
6
7model = SentenceTransformer("all-MiniLM-L6-v2")
8
9# Always compute the vectorstore directory relative to this file
10vectorstore_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../vectorstore"))
11os.makedirs(vectorstore_dir, exist_ok=True)
12
13index_path = os.path.join(vectorstore_dir, "schema_index.faiss")
14meta_path = os.path.join(vectorstore_dir, "schema_meta.pkl")
15
16def build_or_load_index(schema_dict):
17    if os.path.exists(index_path) and os.path.exists(meta_path):
18        print("🔄 Loading FAISS index and metadata from disk...")
19        index = faiss.read_index(index_path)
20        with open(meta_path, "rb") as f:
21            metadata = pickle.load(f)
22        return index, metadata
23
24    print("⚡ Building new FAISS index and metadata (cold start)...")
25    texts, metadata = [], []
26    for table, cols in schema_dict.items():
27        for col in cols:
28            desc = f"{table} - {col}"
29            texts.append(desc)
30            metadata.append({"table": table, "column": col})
31
32    embeddings = model.encode(texts)
33    index = faiss.IndexFlatL2(len(embeddings[0]))
34    index.add(np.array(embeddings).astype("float32"))
35
36    # Directory is already ensured above
37    faiss.write_index(index, index_path)
38    with open(meta_path, "wb") as f:
39        pickle.dump(metadata, f)
40
41    return index, metadata
42