batch embedding

embeddings.py

@@ -4,43 +4,10 @@ import faiss
 import torch
 import numpy as np
 import os
-import json
 
 def log(message):
     print(f"✅ {message}")
 
-
-
-# ✅ Ensure data folder exists
-DATA_DIR = "data"
-os.makedirs(DATA_DIR, exist_ok=True)
-
-# ✅ List of datasets
-datasets_list = {
-    "sales": "goendalf666/sales-conversations",
-    "blended": "blended_skill_talk",
-    "dialog": "daily_dialog",
-    "multiwoz": "multi_woz_v22",
-}
-
-def save_dataset_to_file(dataset_name, dataset):
-    """Save dataset to a local JSON file."""
-    file_path = os.path.join(DATA_DIR, f"{dataset_name}.json")
-    
-    with open(file_path, "w") as f:
-        json.dump(dataset["train"].to_dict(), f)
-    
-    print(f"✅ Saved {dataset_name} to {file_path}")
-
-# ✅ Load & Save all datasets
-for name, dataset_id in datasets_list.items():
-    dataset = load_dataset(dataset_id, split="train")
-    save_dataset_to_file(name, dataset)
-    
-print("✅ All datasets saved locally!")
-
-
-
 # ✅ Load datasets
 datasets = {
     "sales": load_dataset("goendalf666/sales-conversations"),
@@ -49,80 +16,69 @@ datasets = {
     "multiwoz": load_dataset("multi_woz_v22"),
 }
 
-# ✅ Load MiniLM model and tokenizer
-model_name = "sentence-transformers/all-MiniLM-L6-v2"  # Model for embeddings
+# ✅ Load MiniLM model for embeddings
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name)
 
 def embed_text(texts):
+    """Generate embeddings for a batch of texts."""
     inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
     with torch.no_grad():
         embeddings = model(**inputs).last_hidden_state.mean(dim=1).cpu().numpy()
     return embeddings
 
-
-# ✅ Extract and embed the datasets
-def create_embeddings(dataset_name, dataset):
+# ✅ Batch processing function
+def create_embeddings(dataset_name, dataset, batch_size=100):
     print(f"📥 Creating embeddings for {dataset_name}...")
     
     if dataset_name == "goendalf666/sales-conversations":
         texts = [" ".join(row.values()) for row in dataset["train"]]
-
-    elif dataset_name == "AlekseyKorshuk/persona-chat":
-        texts = [" ".join(utterance["candidates"]) for utterance in dataset["train"]["utterances"]]
-
     elif dataset_name == "blended_skill_talk":
         texts = [" ".join(row["free_messages"] + row["guided_messages"]) for row in dataset["train"]]
-
     elif dataset_name == "daily_dialog":
         texts = [" ".join(row["dialog"]) for row in dataset["train"]]
-
     elif dataset_name == "multi_woz_v22":
         texts = [" ".join(row["turns"]["utterance"]) for row in dataset["train"]]
-
     else:
         print(f"⚠️ Warning: Dataset {dataset_name} not handled properly!")
         texts = []
 
-    # ✅ Verify dataset extraction
-    if len(texts) == 0:
-        print(f"❌ ERROR: No text extracted from {dataset_name}! Check dataset structure.")
-    else:
-        print(f"✅ Extracted {len(texts)} texts from {dataset_name}. Sample:\n{texts[:3]}")
-
-    return texts
+    log(f"✅ Extracted {len(texts)} texts from {dataset_name}.")
 
-# ✅ Embed and store in FAISS
-for name, dataset in datasets.items():
-    texts = create_embeddings(name, dataset)
+    # Process in batches
+    all_embeddings = []
+    for i in range(0, len(texts), batch_size):
+        batch = texts[i : i + batch_size]
+        batch_embeddings = embed_text(batch)
+        all_embeddings.append(batch_embeddings)
+        
+        # ✅ Log progress
+        log(f"🚀 Processed {i + len(batch)}/{len(texts)} embeddings for {dataset_name}...")
 
-    if len(texts) > 0:  # ✅ Only embed if texts exist
-        embeddings = embed_text(texts)
-        print(f"✅ Generated embeddings shape: {embeddings.shape}")
+    # Convert list of numpy arrays to a single numpy array
+    all_embeddings = np.vstack(all_embeddings)
+    return all_embeddings
 
-        index = save_embeddings_to_faiss(embeddings)
-        print(f"✅ Embeddings for {name} saved to FAISS.")
+# ✅ Save embeddings to FAISS with unique filename
+def save_embeddings_to_faiss(embeddings, index_name="my_embeddings"):
+    index_file = f"{index_name}.faiss"
+    
+    # ✅ Check if previous FAISS index exists, append if needed
+    if os.path.exists(index_file):
+        log("🔄 Loading existing FAISS index to append...")
+        index = faiss.read_index(index_file)
+        index.add(np.array(embeddings).astype(np.float32))
     else:
-        print(f"⚠️ Skipping embedding for {name} (No valid texts).")
+        index = faiss.IndexFlatL2(embeddings.shape[1])
+        index.add(np.array(embeddings).astype(np.float32))
 
+    faiss.write_index(index, index_file)  # ✅ Save FAISS index
+    log(f"✅ Saved FAISS index: {index_file}")
 
-
-# ✅ Save embeddings to a database
-def save_embeddings_to_faiss(embeddings, index_name="my_embeddings"):
-    print("Saving embeddings to FAISS...")
-    index = faiss.IndexFlatL2(embeddings.shape[1])  # Assuming 512-dimensional embeddings
-    index.add(np.array(embeddings).astype(np.float32))
-    faiss.write_index(index, index_name)  # Save FAISS index to file
-    return index
-
-# ✅ Create embeddings for all datasets
+# ✅ Run embeddings process
 for name, dataset in datasets.items():
-    embeddings = create_embeddings(name, dataset)
-    index = save_embeddings_to_faiss(embeddings)
-    print(f"Embeddings for {name} saved to FAISS.")
-
-
-# ✅ Check FAISS index after saving
-index = faiss.read_index("my_embeddings")  # Load the index
-print(f"📊 FAISS index contains {index.ntotal} vectors.")  # Check how many embeddings were stored
+    embeddings = create_embeddings(name, dataset, batch_size=100)
+    save_embeddings_to_faiss(embeddings, index_name=name)
+    log(f"✅ Embeddings for {name} saved to FAISS.")