update

embeddings.py

@@ -1,19 +1,36 @@
-from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModel
+import os
+import json
 import faiss
 import torch
 import numpy as np
-import os
+from transformers import AutoTokenizer, AutoModel
+
+# ✅ Set up directories
+DATA_DIR = "data"
+os.makedirs(DATA_DIR, exist_ok=True)  # Ensure data directory exists
 
 def log(message):
     print(f"✅ {message}")
 
-# ✅ Load datasets
+# ✅ Load datasets from stored JSON files
+def load_local_dataset(dataset_name):
+    file_path = os.path.join(DATA_DIR, f"{dataset_name}.json")
+    
+    if os.path.exists(file_path):
+        with open(file_path, "r") as f:
+            data = json.load(f)
+        log(f"📂 Loaded {dataset_name} from {file_path}")
+        return data
+    else:
+        log(f"❌ ERROR: {dataset_name} file not found!")
+        return None
+
+# ✅ Load all datasets from storage
 datasets = {
-    "sales": load_dataset("goendalf666/sales-conversations"),
-    "blended": load_dataset("blended_skill_talk"),
-    "dialog": load_dataset("daily_dialog"),
-    "multiwoz": load_dataset("multi_woz_v22"),
+    "sales": load_local_dataset("sales"),
+    "blended": load_local_dataset("blended"),
+    "dialog": load_local_dataset("dialog"),
+    "multiwoz": load_local_dataset("multiwoz"),
 }
 
 # ✅ Load MiniLM model for embeddings
@@ -30,18 +47,19 @@ def embed_text(texts):
 
 # ✅ Batch processing function
 def create_embeddings(dataset_name, dataset, batch_size=100):
-    print(f"📥 Creating embeddings for {dataset_name}...")
+    """Extracts texts, embeds them in batches, and logs progress."""
+    log(f"📥 Creating embeddings for {dataset_name}...")
     
-    if dataset_name == "goendalf666/sales-conversations":
-        texts = [" ".join(row.values()) for row in dataset["train"]]
-    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"]]
+    if dataset_name == "sales":
+        texts = [" ".join(row.values()) for row in dataset]
+    elif dataset_name == "blended":
+        texts = [" ".join(row["free_messages"] + row["guided_messages"]) for row in dataset]
+    elif dataset_name == "dialog":
+        texts = [" ".join(row["dialog"]) for row in dataset]
+    elif dataset_name == "multiwoz":
+        texts = [" ".join(row["turns"]["utterance"]) for row in dataset]
     else:
-        print(f"⚠️ Warning: Dataset {dataset_name} not handled properly!")
+        log(f"⚠️ Warning: Dataset {dataset_name} format unknown!")
         texts = []
 
     log(f"✅ Extracted {len(texts)} texts from {dataset_name}.")
@@ -78,7 +96,8 @@ def save_embeddings_to_faiss(embeddings, index_name="my_embeddings"):
 
 # ✅ Run embeddings process
 for name, dataset in datasets.items():
-    embeddings = create_embeddings(name, dataset, batch_size=100)
-    save_embeddings_to_faiss(embeddings, index_name=name)
-    log(f"✅ Embeddings for {name} saved to FAISS.")
+    if dataset:  # Skip if dataset failed to load
+        embeddings = create_embeddings(name, dataset, batch_size=100)
+        save_embeddings_to_faiss(embeddings, index_name=name)
+        log(f"✅ Embeddings for {name} saved to FAISS.")