| import logging |
| from pathlib import Path |
| from typing import List |
| import pandas as pd |
| import numpy as np |
| import faiss |
| from sentence_transformers import SentenceTransformer |
|
|
| from Data.database.sql_connector import load_preprocessed_projects |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s") |
| logger = logging.getLogger(__name__) |
|
|
| DEFAULT_MODEL = "all-mpnet-base-v2" |
| TEXT_COL = "clean_text" |
| TITLE_COL = "project_title" |
| TECH_COL = "technologies" |
|
|
| _PROJECT_ROOT = Path(__file__).resolve().parents[2] |
| MODEL_DIR = _PROJECT_ROOT / "models" |
| INDEX_PATH = MODEL_DIR / "faiss_index.bin" |
| META_PATH = MODEL_DIR / "metadata.parquet" |
|
|
| class ProjectEmbedder: |
| def __init__(self, model_name: str = DEFAULT_MODEL): |
| logger.info(f"Loading embedding model: {model_name}") |
| self.model = SentenceTransformer(model_name) |
| self.index = None |
| self.metadata = None |
|
|
| def generate_embeddings(self, texts: List[str], batch_size: int = 64) -> np.ndarray: |
| logger.info(f"Generating embeddings for {len(texts)} projects...") |
| vectors = self.model.encode( |
| texts, |
| batch_size=batch_size, |
| show_progress_bar=True, |
| convert_to_numpy=True, |
| normalize_embeddings=True |
| ) |
| return vectors.astype("float32") |
|
|
| def build_index(self, df: pd.DataFrame): |
| """Build FAISS cosine index.""" |
| self.metadata = df.copy() |
| self.metadata = self.metadata.reset_index(drop=True) |
|
|
| for col in [TITLE_COL, TEXT_COL]: |
| if col not in self.metadata.columns: |
| self.metadata[col] = "" |
|
|
| if TECH_COL not in self.metadata.columns: |
| self.metadata[TECH_COL] = "" |
|
|
| FEATURE_COL = "features" |
| if FEATURE_COL not in self.metadata.columns: |
| self.metadata[FEATURE_COL] = "" |
|
|
| feature_text = self.metadata[FEATURE_COL].fillna("").astype(str) |
| rich_texts = ( |
| self.metadata[TITLE_COL].fillna("").astype(str) |
| + " " |
| + self.metadata[TEXT_COL].fillna("").astype(str) |
| + " " |
| + feature_text |
| ).tolist() |
|
|
| embeddings = self.generate_embeddings(rich_texts) |
| self.embeddings = embeddings |
| dim = embeddings.shape[1] |
| |
| base_index = faiss.IndexFlatIP(dim) |
| self.index = faiss.IndexIDMap(base_index) |
| ids = np.arange(len(self.metadata)).astype("int64") |
| |
| self.index.add_with_ids(embeddings, ids) |
| logger.info(f"FAISS index built successfully with {self.index.ntotal} vectors.") |
|
|
| def save_artifacts(self, folder: str = "models"): |
| path = Path(folder) |
| path.mkdir(parents=True, exist_ok=True) |
| faiss.write_index(self.index, str(path / "faiss_index.bin")) |
| self.metadata.to_parquet(path / "metadata.parquet", index=False) |
| if hasattr(self, "embeddings"): |
| np.save(str(path / "project_embeddings.npy"), self.embeddings) |
| logger.info(f"Artifacts saved to {folder}") |
|
|
| def load_artifacts(self, folder: str = "models"): |
| path = Path(folder) |
| self.index = faiss.read_index(str(path / "faiss_index.bin")) |
| self.metadata = pd.read_parquet(path / "metadata.parquet") |
| logger.info("Artifacts loaded successfully.") |
|
|
| def train_embedding_engine(): |
| logger.info("Loading processed dataset from Azure SQL...") |
| df = load_preprocessed_projects() |
| |
| engine = ProjectEmbedder() |
| engine.build_index(df) |
| engine.save_artifacts(str(MODEL_DIR)) |
| |
| logger.info("Embedding engine completed successfully.") |
| return engine |
|
|
