Create precompute_cap_embeddings.py

precompute_cap_embeddings.py

@@ -0,0 +1,69 @@
+import os
+import logging
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from openai import OpenAI
+from scipy.sparse import save_npz
+import pickle
+from datasets import load_from_disk
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+# === Logging setup ===
+logger = logging.getLogger("precompute")
+logging.basicConfig(level=logging.INFO)
+
+# === API keys ===
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+openai_client = OpenAI(api_key=OPENAI_API_KEY)
+
+# === Load CAP dataset ===
+LOCAL_PATH = "/data/cap_dataset"
+cap_dataset = load_from_disk(LOCAL_PATH)
+cap_texts = [doc['text'] for doc in cap_dataset]
+logger.info(f"Loaded {len(cap_texts)} CAP texts.")
+
+# === TF-IDF Precomputation ===
+if not (os.path.exists("/data/cap_tfidf.pkl") and os.path.exists("/data/cap_tfidf_matrix.npz")):
+    logger.info("Precomputing TF-IDF...")
+    tfidf = TfidfVectorizer(stop_words='english', max_features=100_000)
+    tfidf_matrix = tfidf.fit_transform(cap_texts)
+    with open("/data/cap_tfidf.pkl", 'wb') as f:
+        pickle.dump(tfidf, f)
+    save_npz("/data/cap_tfidf_matrix.npz", tfidf_matrix)
+    logger.info("✅ Saved TF-IDF cache files.")
+else:
+    logger.info("TF-IDF cache files already exist, skipping.")
+
+# === GTE Embeddings Precomputation ===
+if not os.path.exists("/data/cap_gte.npy"):
+    logger.info("Precomputing GTE embeddings...")
+    encoder_gte = SentenceTransformer("Alibaba-NLP/gte-Qwen2-1.5B-instruct")
+    embeddings_gte = encoder_gte.encode(cap_texts, normalize_embeddings=True)
+    np.save("/data/cap_gte.npy", embeddings_gte)
+    logger.info("✅ Saved GTE embeddings.")
+else:
+    logger.info("GTE embeddings cache file already exists, skipping.")
+
+# === OpenAI Embeddings Precomputation ===
+if not os.path.exists("/data/cap_openai.npy"):
+    logger.info("Precomputing OpenAI embeddings...")
+    def get_openai_embeddings(texts):
+        chunk_size = 100  # Adjust based on average text length and token limit
+        embeddings = []
+        for i in range(0, len(texts), chunk_size):
+            chunk = texts[i:i + chunk_size]
+            response = openai_client.embeddings.create(
+                model="text-embedding-3-large",
+                input=chunk
+            )
+            embeddings.extend([item.embedding for item in response.data])
+            logger.info(f"Processed chunk {i//chunk_size + 1} of {len(texts)//chunk_size + 1}")
+            time.sleep(1)  # Rate limit buffer for Tier 5
+        return np.array(embeddings)
+    embeddings_openai = get_openai_embeddings(cap_texts)
+    np.save("/data/cap_openai.npy", embeddings_openai)
+    logger.info("✅ Saved OpenAI embeddings.")
+else:
+    logger.info("OpenAI embeddings cache file already exists, skipping.")
+
+logger.info("✅ Precomputation completed. Cache files are ready for use.")