Update precompute_cap_embeddings.py

precompute_cap_embeddings.py

@@ -12,40 +12,48 @@ from sklearn.feature_extraction.text import TfidfVectorizer
 logger = logging.getLogger("precompute")
 logging.basicConfig(level=logging.INFO)
 
-# === API keys ===
+# === API key handling ===
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+if not OPENAI_API_KEY:
+    OPENAI_API_KEY = input("Please enter your OpenAI API Key (set OPENAI_API_KEY environment variable for future runs): ")
+    if not OPENAI_API_KEY:
+        raise EnvironmentError("OPENAI_API_KEY must be provided either as an environment variable or input.")
 openai_client = OpenAI(api_key=OPENAI_API_KEY)
 
 # === Load CAP dataset ===
-LOCAL_PATH = "/data/cap_dataset"
+LOCAL_PATH = "./cap_dataset"  # Local path for testing
+if not os.path.exists(LOCAL_PATH):
+    raise FileNotFoundError(f"CAP dataset not found at {LOCAL_PATH}. Download it first.")
 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")):
+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:
+    os.makedirs("./data", exist_ok=True)
+    with open("./data/cap_tfidf.pkl", 'wb') as f:
         pickle.dump(tfidf, f)
-    save_npz("/data/cap_tfidf_matrix.npz", tfidf_matrix)
+    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"):
+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)
+    os.makedirs("./data", exist_ok=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"):
+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
@@ -61,7 +69,8 @@ if not os.path.exists("/data/cap_openai.npy"):
             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)
+    os.makedirs("./data", exist_ok=True)
+    np.save("./data/cap_openai.npy", embeddings_openai)
     logger.info("✅ Saved OpenAI embeddings.")
 else:
     logger.info("OpenAI embeddings cache file already exists, skipping.")