Spaces:
Sleeping
Sleeping
| """ | |
| scripts/build_features.py | |
| -------------------------- | |
| Parse a board game rulebook into semantic chunks, embed them with | |
| sentence-transformers, and store a FAISS index + chunk metadata. | |
| Run once per game before using the RAG pipeline (from project root): | |
| .venv/bin/python scripts/build_features.py catan | |
| .venv/bin/python scripts/build_features.py monopoly | |
| """ | |
| import re | |
| import json | |
| import sys | |
| from pathlib import Path | |
| # --------------------------------------------------------------------------- | |
| # Shared utility | |
| # --------------------------------------------------------------------------- | |
| def _position_based_chunks(text: str, sections: list[tuple], prefix: str) -> list[dict]: | |
| """ | |
| Generic position-based chunker. | |
| Finds each section header via regex, then slices the text between headers. | |
| """ | |
| positions = [] | |
| for title, pattern in sections: | |
| m = re.search(pattern, text) | |
| if m: | |
| positions.append((m.start(), title)) | |
| else: | |
| print(f" WARNING: could not locate '{title}'", file=sys.stderr) | |
| positions.sort(key=lambda x: x[0]) | |
| chunks = [] | |
| for i, (start, title) in enumerate(positions): | |
| end = positions[i + 1][0] if i + 1 < len(positions) else len(text) | |
| content = text[start:end].strip() | |
| if content: | |
| chunks.append({"title": f"{prefix}: {title}", "text": content, "source": prefix}) | |
| return chunks | |
| # --------------------------------------------------------------------------- | |
| # Catan chunkers | |
| # --------------------------------------------------------------------------- | |
| _CATAN_GAME_RULES_SECTIONS = [ | |
| ("GAME COMPONENTS", r"GAME COMPONENTS"), | |
| ("ISLAND CONSTRUCTION AND SETUP", r"CONSTRUCTING THE ISLAND"), | |
| ("TURN OVERVIEW", r"TURN OVERVIEW"), | |
| ("RESOURCE PRODUCTION", r"1\. RESOURCE PRODUCTION"), | |
| ("TRADE (Domestic and Maritime)", r"2\. TRADE"), | |
| ("BUILD (Roads, Settlements, Cities, Development Cards)", r"3\. BUILD"), | |
| ('ROLLING A "7" AND ACTIVATING THE ROBBER', r'a\) Rolling a'), | |
| ("PLAYING DEVELOPMENT CARDS (Knight, Progress, Victory Point)", r"b\) Playing Development Cards"), | |
| ("ENDING THE GAME", r"ENDING THE GAME"), | |
| ] | |
| def _is_almanac_header(line: str) -> bool: | |
| """Return True if the line looks like an ALL-CAPS Almanac entry header.""" | |
| line = line.strip() | |
| if not line or "=" in line or "Pages" in line: | |
| return False | |
| if len(line) < 3 or len(line) > 65: | |
| return False | |
| upper = sum(1 for c in line if c.isupper()) | |
| lower = sum(1 for c in line if c.islower()) | |
| total_alpha = upper + lower | |
| if total_alpha < 3: | |
| return False | |
| return upper / total_alpha >= 0.80 | |
| def _parse_catan_almanac(text: str) -> list[dict]: | |
| """Split the Catan Almanac section into one chunk per entry.""" | |
| chunks = [] | |
| current_title = None | |
| current_lines: list[str] = [] | |
| for line in text.split("\n"): | |
| stripped = line.strip() | |
| if stripped and _is_almanac_header(stripped): | |
| if current_title and current_lines: | |
| content = "\n".join(current_lines).strip() | |
| if len(content) > 30: | |
| chunks.append({"title": f"Almanac: {current_title}", "text": content, "source": "Almanac"}) | |
| current_title = stripped | |
| current_lines = [stripped] | |
| else: | |
| if current_title is not None: | |
| current_lines.append(line) | |
| if current_title and current_lines: | |
| content = "\n".join(current_lines).strip() | |
| if len(content) > 30: | |
| chunks.append({"title": f"Almanac: {current_title}", "text": content, "source": "Almanac"}) | |
| return chunks | |
| def chunk_catan(filepath: str) -> list[dict]: | |
| """Parse the full Catan rulebook into semantic chunks.""" | |
| with open(filepath) as f: | |
| text = f.read() | |
| almanac_marker = "ALMANAC (Pages 6-15)" | |
| if almanac_marker not in text: | |
| print("WARNING: Could not find ALMANAC marker; treating entire file as Game Rules.", file=sys.stderr) | |
| return _position_based_chunks(text, _CATAN_GAME_RULES_SECTIONS, "Game Rules") | |
| split_pos = text.index(almanac_marker) | |
| chunks = _position_based_chunks(text[:split_pos], _CATAN_GAME_RULES_SECTIONS, "Game Rules") | |
| chunks += _parse_catan_almanac(text[split_pos:]) | |
| return chunks | |
| # --------------------------------------------------------------------------- | |
| # Monopoly chunker | |
| # --------------------------------------------------------------------------- | |
| _MONOPOLY_SECTIONS = [ | |
| ("OVERVIEW", r"MONOPOLY - Official Rules"), | |
| ("SPEED DIE RULES", r"SPEED DIE RULES"), | |
| ("OBJECT", r"OBJECT:"), | |
| ("PREPARATION", r"PREPARATION:"), | |
| ("THE BANKER", r"BANKER:"), | |
| ("THE BANK", r"THE BANK:"), | |
| ("THE PLAY", r"THE PLAY:"), | |
| ("GO", r'"GO":'), | |
| ("BUYING PROPERTY", r"BUYING PROPERTY:"), | |
| ("PAYING RENT", r"PAYING RENT:"), | |
| ("CHANCE AND COMMUNITY CHEST", r'"CHANCE" AND'), | |
| ("INCOME TAX", r'"INCOME TAX":'), | |
| ("JAIL", r'"JAIL":'), | |
| ("FREE PARKING", r'"FREE PARKING":'), | |
| ("HOUSES", r"HOUSES:"), | |
| ("HOTELS", r"HOTELS:"), | |
| ("BUILDING SHORTAGES", r"BUILDING SHORTAGES:"), | |
| ("SELLING PROPERTY", r"SELLING PROPERTY:"), | |
| ("MORTGAGES", r"MORTGAGES:"), | |
| ("BANKRUPTCY", r"BANKRUPTCY:"), | |
| ("MISCELLANEOUS", r"MISCELLANEOUS:"), | |
| ] | |
| def chunk_monopoly(filepath: str) -> list[dict]: | |
| """Parse the Monopoly rulebook into semantic chunks.""" | |
| with open(filepath) as f: | |
| text = f.read() | |
| return _position_based_chunks(text, _MONOPOLY_SECTIONS, "Monopoly") | |
| # --------------------------------------------------------------------------- | |
| # Game registry — add new games here | |
| # --------------------------------------------------------------------------- | |
| GAMES: dict[str, callable] = { | |
| "catan": chunk_catan, | |
| "monopoly": chunk_monopoly, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Embedding + FAISS index (deep learning retriever) | |
| # --------------------------------------------------------------------------- | |
| def build_faiss_index(chunks: list[dict]): | |
| """Embed all chunks and build a cosine-similarity FAISS index.""" | |
| try: | |
| import faiss | |
| from sentence_transformers import SentenceTransformer | |
| except ImportError as e: | |
| print(f"Missing dependency: {e}\nRun: .venv/bin/pip install -r requirements.txt") | |
| sys.exit(1) | |
| print("Loading sentence-transformers model (all-MiniLM-L6-v2)...") | |
| model = SentenceTransformer("all-MiniLM-L6-v2") | |
| texts = [c["text"] for c in chunks] | |
| print(f"Embedding {len(texts)} chunks...") | |
| embeddings = model.encode(texts, show_progress_bar=True, convert_to_numpy=True) | |
| faiss.normalize_L2(embeddings) | |
| dim = embeddings.shape[1] | |
| index = faiss.IndexFlatIP(dim) | |
| index.add(embeddings) | |
| print(f"FAISS index: {index.ntotal} vectors, dim={dim}") | |
| return index | |
| # --------------------------------------------------------------------------- | |
| # TF-IDF index (classical ML retriever) | |
| # --------------------------------------------------------------------------- | |
| def build_tfidf_index(chunks: list[dict], out_dir: Path): | |
| """Fit a TF-IDF vectorizer on all chunks and save the matrix + vectorizer.""" | |
| try: | |
| import pickle | |
| import scipy.sparse | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| except ImportError as e: | |
| print(f"Missing dependency: {e}\nRun: .venv/bin/pip install -r requirements.txt") | |
| sys.exit(1) | |
| texts = [c["text"] for c in chunks] | |
| print(f"Building TF-IDF index over {len(texts)} chunks...") | |
| vectorizer = TfidfVectorizer(stop_words="english", ngram_range=(1, 2)) | |
| matrix = vectorizer.fit_transform(texts) | |
| print(f"TF-IDF matrix: {matrix.shape[0]} docs × {matrix.shape[1]} terms") | |
| scipy.sparse.save_npz(str(out_dir / "tfidf_matrix.npz"), matrix) | |
| with open(str(out_dir / "tfidf_vectorizer.pkl"), "wb") as f: | |
| pickle.dump(vectorizer, f) | |
| print(f"Saved TF-IDF artifacts -> '{out_dir}'") | |
| return vectorizer, matrix | |
| # --------------------------------------------------------------------------- | |
| # Main | |
| # --------------------------------------------------------------------------- | |
| def main(): | |
| import argparse | |
| parser = argparse.ArgumentParser(description="Build FAISS index for a board game rulebook.") | |
| parser.add_argument("game", choices=list(GAMES.keys()), help="Game to index") | |
| args = parser.parse_args() | |
| game = args.game | |
| rulebook_path = f"data/raw/{game}_rulebook.txt" | |
| out_dir = Path(f"models/{game}") | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| chunks_path = str(out_dir / "chunks.json") | |
| index_path = str(out_dir / "faiss.index") | |
| print(f"Reading rulebook from '{rulebook_path}'...") | |
| chunks = GAMES[game](rulebook_path) | |
| print(f"\nCreated {len(chunks)} chunks:\n") | |
| for i, c in enumerate(chunks): | |
| print(f" {i+1:2d}. {c['title']:<70} ({len(c['text'])} chars)") | |
| with open(chunks_path, "w") as f: | |
| json.dump(chunks, f, indent=2, ensure_ascii=False) | |
| print(f"\nSaved chunks -> '{chunks_path}'") | |
| import faiss | |
| index = build_faiss_index(chunks) | |
| faiss.write_index(index, index_path) | |
| print(f"Saved index -> '{index_path}'") | |
| build_tfidf_index(chunks, out_dir) | |
| print(f"\nDone. You can now run: .venv/bin/python scripts/model.py {game} 'your question'") | |
| if __name__ == "__main__": | |
| main() | |