Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| import argparse | |
| import sys | |
| import string | |
| import math | |
| from nltk.stem import PorterStemmer | |
| from lib.inverted_index import InvertedIndex | |
| from lib.search_utils import load_stopwords | |
| from lib.keyword_search import bm25_idf_command, bm25_tf_command | |
| from lib.constants import BM25_K1, BM25_B | |
| def normalize_to_tokens(text: str) -> list[str]: | |
| translator = str.maketrans("", "", string.punctuation) | |
| stemmer = PorterStemmer() | |
| stopwords = set(load_stopwords()) | |
| clean = text.translate(translator).lower() | |
| tokens = [t for t in clean.split() if t] | |
| tokens = [t for t in tokens if t not in stopwords] | |
| tokens = [stemmer.stem(t) for t in tokens] | |
| return tokens | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description="Keyword Search CLI") | |
| subparsers = parser.add_subparsers(dest="command", help="Available commands") | |
| search_parser = subparsers.add_parser("search", help="Search movies using inverted index") | |
| search_parser.add_argument("query", type=str, help="Search query") | |
| subparsers.add_parser("build", help="Build inverted index") | |
| tf_parser = subparsers.add_parser("tf", help="Get term frequency for a term in a document") | |
| tf_parser.add_argument("doc_id", type=int, help="Document ID") | |
| tf_parser.add_argument("term", type=str, help="Term to search for") | |
| idf_parser = subparsers.add_parser("idf", help="Get inverse document frequency for a term") | |
| idf_parser.add_argument("term", type=str, help="Term to calculate IDF for") | |
| tfidf_parser = subparsers.add_parser("tfidf", help="Get TF-IDF score for a term in a document") | |
| tfidf_parser.add_argument("doc_id", type=int, help="Document ID") | |
| tfidf_parser.add_argument("term", type=str, help="Term to calculate TF-IDF for") | |
| bm25_idf_parser = subparsers.add_parser("bm25idf", help="Get BM25 IDF score for a given term") | |
| bm25_idf_parser.add_argument("term", type=str, help="Term to get BM25 IDF score for") | |
| bm25_tf_parser = subparsers.add_parser("bm25tf", help="Get BM25 TF score for a given document ID and term") | |
| bm25_tf_parser.add_argument("doc_id", type=int, help="Document ID") | |
| bm25_tf_parser.add_argument("term", type=str, help="Term to get BM25 TF score for") | |
| bm25_tf_parser.add_argument("k1", type=float, nargs='?', default=BM25_K1, help="Tunable BM25 K1 parameter") | |
| bm25_tf_parser.add_argument("b", type=float, nargs='?', default=BM25_B, help="Tunable BM25 b parameter") | |
| bm25search_parser = subparsers.add_parser("bm25search", help="Search movies using full BM25 scoring") | |
| bm25search_parser.add_argument("query", type=str, help="Search query") | |
| bm25search_parser.add_argument("--limit", type=int, default=5, help="Number of results to return") | |
| args = parser.parse_args() | |
| match args.command: | |
| case "build": | |
| idx = InvertedIndex() | |
| idx.build() | |
| idx.save() | |
| print("Built and saved inverted index.") | |
| case "search": | |
| print(f"Searching for: {args.query}") | |
| idx = InvertedIndex() | |
| try: | |
| idx.load() | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| query_tokens = normalize_to_tokens(args.query) | |
| results: list[int] = [] | |
| seen = set() | |
| for token in query_tokens: | |
| doc_ids = idx.get_documents(token) | |
| for doc_id in doc_ids: | |
| if doc_id not in seen: | |
| seen.add(doc_id) | |
| results.append(doc_id) | |
| if len(results) >= 5: | |
| break | |
| if len(results) >= 5: | |
| break | |
| for i, doc_id in enumerate(results, 1): | |
| movie = idx.docmap[doc_id] | |
| print(f"{i}. {movie['title']} ({doc_id})") | |
| case "tf": | |
| idx = InvertedIndex() | |
| try: | |
| idx.load() | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| tf = idx.get_tf(args.doc_id, args.term) | |
| print(tf) | |
| case "idf": | |
| idx = InvertedIndex() | |
| try: | |
| idx.load() | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| tokens = normalize_to_tokens(args.term) | |
| if len(tokens) == 0: | |
| print(f"Inverse document frequency of '{args.term}': 0.00") | |
| sys.exit(0) | |
| if len(tokens) > 1: | |
| print("Error: idf expects a single token term") | |
| sys.exit(1) | |
| token = tokens[0] | |
| # Calculate IDF: log((N + 1) / (n_t + 1)) | |
| N = len(idx.docmap) | |
| doc_ids = idx.get_documents(token) | |
| n_t = len(doc_ids) | |
| idf = math.log((N + 1) / (n_t + 1)) | |
| print(f"Inverse document frequency of '{args.term}': {idf:.2f}") | |
| case "tfidf": | |
| idx = InvertedIndex() | |
| try: | |
| idx.load() | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| tokens = normalize_to_tokens(args.term) | |
| if len(tokens) == 0: | |
| print(f"TF-IDF score of '{args.term}' in document '{args.doc_id}': 0.00") | |
| sys.exit(0) | |
| if len(tokens) > 1: | |
| print("Error: tfidf expects a single token term") | |
| sys.exit(1) | |
| token = tokens[0] | |
| # Calculate TF | |
| tf = idx.get_tf(args.doc_id, token) | |
| # Calculate IDF | |
| N = len(idx.docmap) | |
| doc_ids = idx.get_documents(token) | |
| n_t = len(doc_ids) | |
| idf = math.log((N + 1) / (n_t + 1)) | |
| # Calculate TF-IDF | |
| tf_idf = tf * idf | |
| print(f"TF-IDF score of '{args.term}' in document '{args.doc_id}': {tf_idf:.2f}") | |
| case "bm25idf": | |
| try: | |
| bm25idf = bm25_idf_command(args.term) | |
| print(f"BM25 IDF score of '{args.term}': {bm25idf:.2f}") | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| case "bm25tf": | |
| try: | |
| bm25tf = bm25_tf_command(args.doc_id, args.term, args.k1, args.b) | |
| print(f"BM25 TF score of '{args.term}' in document '{args.doc_id}': {bm25tf:.2f}") | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| case "bm25search": | |
| try: | |
| idx = InvertedIndex() | |
| idx.load() | |
| limit = args.limit if hasattr(args, 'limit') else 5 | |
| results = idx.bm25_search(args.query, limit) | |
| for i, (doc_id, score) in enumerate(results, 1): | |
| movie = idx.docmap[doc_id] | |
| print(f"{i}. ({doc_id}) {movie['title']} - Score: {score:.2f}") | |
| except FileNotFoundError: | |
| print("Error: index not found. Run: uv run cli/keyword_search_cli.py build") | |
| sys.exit(1) | |
| case _: | |
| parser.print_help() | |
| if __name__ == "__main__": | |
| main() | |