import pandas as pd
import faiss
import pickle
from sentence_transformers import SentenceTransformer

print("Loading dataset...")

df = pd.read_csv("shl_individual_tests_catalog.csv")

# Combine useful columns into one searchable text
df["combined_text"] = (

    "Assessment: " +
    df["name"].fillna("").astype(str) +

    ". Description: " +
    df["description"].fillna("").astype(str) +

    ". Test Type: " +
    df["test_type"].fillna("").astype(str) +

    ". Skills assessment hiring recruitment coding personality cognitive communication developer"
)

print("Loading embedding model...")

model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2"
)

print("Generating embeddings...")

embeddings = model.encode(
    df["combined_text"].tolist(),
    show_progress_bar=True
)

print("Creating FAISS index...")

dimension = embeddings.shape[1]

index = faiss.IndexFlatL2(dimension)

index.add(embeddings)

# Save index
faiss.write_index(
    index,
    "shl_index.faiss"
)

# Save catalog separately
pickle.dump(
    df,
    open(
        "catalog.pkl",
        "wb"
    )
)

print("Done")
print(f"Indexed {len(df)} assessments")