shl-assignment / embeddings.py
poushali22's picture
Final clean deployment
47d8d24
Raw
History Blame Contribute Delete
1.15 kB
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")