Upload src/save_doc_embeddings_to_redis.py with huggingface_hub

src/save_doc_embeddings_to_redis.py

@@ -0,0 +1,96 @@
+import torch
+import torch.nn as nn
+import pickle
+import glob
+import json
+import redis
+import numpy as np
+from tqdm import tqdm
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+# Redis Cloud connection (replace with your actual credentials or use environment variables)
+REDIS_HOST = os.environ.get('REDIS_HOST', 'your-redis-host')
+REDIS_PORT = int(os.environ.get('REDIS_PORT', 12345))
+REDIS_PASSWORD = os.environ.get('REDIS_PASSWORD', 'your-redis-password')
+r = redis.Redis(
+    host=REDIS_HOST,
+    port=REDIS_PORT,
+    password=REDIS_PASSWORD,
+    decode_responses=False  # binary-safe
+)
+
+# Load tokenizer
+with open('cbow/tkn_words_to_ids.pkl', 'rb') as f:
+    words_to_ids = pickle.load(f)
+vocab_size = len(words_to_ids)
+embedding_dim = 128  # Change if needed
+
+# Load latest CBOW checkpoint for embedding layer
+checkpoint_files = glob.glob('cbow/checkpoints/*.pth')
+latest_checkpoint = max(checkpoint_files, key=os.path.getctime)
+state_dict = torch.load(latest_checkpoint, map_location='cpu')
+
+embedding_layer = nn.Embedding(vocab_size, embedding_dim)
+embedding_layer.weight.data.copy_(state_dict['emb.weight'])
+embedding_layer.weight.requires_grad = False
+
+# Define DocTower (copy from simple_dual_encoder_rnn.py)
+class DocTower(nn.Module):
+    def __init__(self, embedding_layer, hidden_size):
+        super().__init__()
+        self.embedding = embedding_layer
+        self.embedding.weight.requires_grad = False
+        self.rnn = nn.GRU(input_size=self.embedding.embedding_dim, hidden_size=hidden_size, batch_first=True)
+
+    def forward(self, x):
+        if not x:
+            return None
+        x = torch.tensor(x, dtype=torch.long).unsqueeze(0)
+        embeds = self.embedding(x)
+        _, h_n = self.rnn(embeds)
+        return h_n.squeeze(0).squeeze(0)
+
+# Load doc tower weights (if you saved them separately, load here)
+# Otherwise, use the same initialization as in training
+hidden_size = 128  # Set to your trained hidden size
+
+docTower = DocTower(embedding_layer, hidden_size)
+# Optionally: docTower.load_state_dict(torch.load('doc_tower.pth'))
+docTower.eval()
+
+# Load tokenized triples
+with open('tokenized_triples.json', 'r') as f:
+    triples_data = json.load(f)
+
+# Collect all unique positive documents
+seen = set()
+documents = []
+for split in ['train', 'validation', 'test']:
+    for triple in triples_data[split]:
+        doc_text = triple['positive_document']
+        doc_tokens = tuple(triple['positive_document_tokens'])
+        if doc_tokens not in seen:
+            seen.add(doc_tokens)
+            documents.append((doc_tokens, doc_text))
+
+print(f"Found {len(documents)} unique positive documents.")
+
+def save_doc_embedding_to_redis(doc_id, embedding, text):
+    r.hset(doc_id, mapping={
+        'embedding': embedding.astype(np.float32).tobytes(),
+        'text': text,
+        'doc_id': doc_id
+    })
+
+# Compute and save embeddings
+docTower.eval()
+with torch.no_grad():
+    for idx, (doc_tokens, doc_text) in enumerate(tqdm(documents, desc='Saving doc embeddings to Redis')):
+        embedding = docTower(list(doc_tokens)).detach().cpu().numpy()
+        doc_id = f"doc:{idx}"
+        save_doc_embedding_to_redis(doc_id, embedding, doc_text)
+
+print(f"Saved {len(documents)} doc embeddings to Redis Cloud.")