Team Ai
Apppublic

Decoder2704/python-chatbot-jay

sourceHugging Faceupdated 7mo agoView on Hugging Face
0likes
build_retrieval_artifacts.py149 linesDownload Raw Back to scripts
1#!/usr/bin/env python32"""3Build TF-IDF retrieval artifacts from final_chatbot_data.csv.4 5Parity with Chatbot-V5.ipynb (cells 14–15): TfidfVectorizer fit on concatenated6Question and Answer columns; transform Questions only.7 8Default input:  backend/data/final_chatbot_data.csv9Default outputs: backend/artifacts/vectorizer.pkl, question_vectors.pkl, chatbot_df.pkl10 11Paths resolve from this script file location (not the process working directory).12"""13 14from __future__ import annotations15 16import argparse17import logging18import sys19from pathlib import Path20 21import joblib22import numpy as np23import pandas as pd24from sklearn.feature_extraction.text import TfidfVectorizer25 26_BACKEND_DIR = Path(__file__).resolve().parent.parent27if str(_BACKEND_DIR) not in sys.path:28    sys.path.insert(0, str(_BACKEND_DIR))29 30from app.paths import ARTIFACTS_DIR, DATA_DIR, ensure_data_and_artifacts_dirs31 32DEFAULT_INPUT_CSV = DATA_DIR / "final_chatbot_data.csv"33DEFAULT_ARTIFACTS_DIR = ARTIFACTS_DIR34 35VECTORIZER_NAME = "vectorizer.pkl"36QUESTION_VECTORS_NAME = "question_vectors.pkl"37CHATBOT_DF_NAME = "chatbot_df.pkl"38 39 40logging.basicConfig(41    level=logging.INFO,42    format="%(levelname)s: %(message)s",43)44logger = logging.getLogger(__name__)45 46 47def load_final_dataset(path: Path) -> pd.DataFrame:48    logger.info("Loading dataset from %s", path)49    df = pd.read_csv(path)50    logger.info("Loaded %s rows, %s columns: %s", f"{len(df):,}", len(df.columns), list(df.columns))51    return df52 53 54def build_retrieval_artifacts(df: pd.DataFrame) -> tuple[TfidfVectorizer, object]:55    """Same as Chatbot-V5.ipynb cells 14–15."""56    vectorizer = TfidfVectorizer()57    vectorizer.fit(np.concatenate((df.Question, df.Answer)))58    Question_vectors = vectorizer.transform(df.Question)59    return vectorizer, Question_vectors60 61 62def save_artifacts(63    vectorizer: TfidfVectorizer,64    Question_vectors: object,65    df: pd.DataFrame,66    artifacts_dir: Path,67) -> None:68    artifacts_dir.mkdir(parents=True, exist_ok=True)69 70    v_path = artifacts_dir / VECTORIZER_NAME71    qv_path = artifacts_dir / QUESTION_VECTORS_NAME72    df_path = artifacts_dir / CHATBOT_DF_NAME73 74    logger.info("Writing %s", v_path)75    joblib.dump(vectorizer, v_path)76 77    logger.info("Writing %s (shape %s)", qv_path, getattr(Question_vectors, "shape", "?"))78    joblib.dump(Question_vectors, qv_path)79 80    logger.info("Writing %s", df_path)81    joblib.dump(df, df_path)82 83 84def parse_args() -> argparse.Namespace:85    p = argparse.ArgumentParser(86        description="Build TF-IDF artifacts (Chatbot-V5.ipynb retrieval step).",87    )88    p.add_argument(89        "--input-csv",90        type=Path,91        default=DEFAULT_INPUT_CSV,92        help=f"Path to final_chatbot_data.csv (default: {DEFAULT_INPUT_CSV})",93    )94    p.add_argument(95        "--artifacts-dir",96        type=Path,97        default=DEFAULT_ARTIFACTS_DIR,98        help=f"Directory for .pkl outputs (default: {DEFAULT_ARTIFACTS_DIR})",99    )100    return p.parse_args()101 102 103def main() -> int:104    ensure_data_and_artifacts_dirs()105    args = parse_args()106    input_csv = args.input_csv.resolve()107    artifacts_dir = args.artifacts_dir.resolve()108 109    v_path = artifacts_dir / VECTORIZER_NAME110    qv_path = artifacts_dir / QUESTION_VECTORS_NAME111    df_path = artifacts_dir / CHATBOT_DF_NAME112 113    if not input_csv.is_file():114        print(f"ERROR: input CSV not found: {input_csv}", file=sys.stderr)115        return 1116 117    try:118        df = load_final_dataset(input_csv)119        vectorizer, Question_vectors = build_retrieval_artifacts(df)120        save_artifacts(vectorizer, Question_vectors, df, artifacts_dir)121    except KeyError as exc:122        print(123            f"ERROR: dataset must include Question and Answer columns ({exc})",124            file=sys.stderr,125        )126        return 1127    except UnicodeDecodeError as exc:128        print(f"ERROR: failed to decode CSV: {exc}", file=sys.stderr)129        return 1130    except OSError as exc:131        print(f"ERROR: file read/write failed: {exc}", file=sys.stderr)132        return 1133    except Exception as exc:134        logger.exception("Failed to build retrieval artifacts")135        print(f"ERROR: {exc}", file=sys.stderr)136        return 1137 138    print(f"input:              {input_csv}")139    print(f"output vectorizer:  {v_path}")140    print(f"output vectors:     {qv_path}")141    print(f"output dataframe:   {df_path}")142    print(f"dataframe shape:    {df.shape}")143    print(f"vectors shape:      {Question_vectors.shape}")144    return 0145 146 147if __name__ == "__main__":148    raise SystemExit(main())149