restoring datasets

app.py

@@ -1,78 +1,83 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModel
 import faiss
+import torch
 import numpy as np
-import os
 import time
-import threading  # ✅ Run embeddings in parallel
-
-# ✅ Ensure FAISS is installed
-os.system("pip install faiss-cpu")
 
 def log(message):
     print(f"✅ {message}")
 
-# ✅ Step 1: Run Embeddings in a Separate Thread
-def run_embeddings():
-    log("🚀 Running embeddings script in background...")
-    import embeddings  # ✅ This will automatically run embeddings.py
-    log("✅ Embeddings process finished.")
+# ✅ Load datasets dynamically
+datasets = {
+    "sales": load_dataset("goendalf666/sales-conversations"),
+    "blended": load_dataset("blended_skill_talk"),
+    "dialog": load_dataset("daily_dialog"),
+    "multiwoz": load_dataset("multi_woz_v22"),
+}
 
-embedding_thread = threading.Thread(target=run_embeddings)
-embedding_thread.start()  # ✅ Start embedding in background
+# ✅ 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)
 
-# ✅ Step 2: Check FAISS index
-def check_faiss():
-    index_path = "my_embeddings.faiss"  # Ensure file has .faiss extension
+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
 
-    if not os.path.exists(index_path):
-        return "⚠️ No FAISS index found! Embeddings might still be processing."
+# ✅ Batch processing function
+def create_embeddings(dataset_name, dataset, batch_size=100):
+    log(f"📥 Creating embeddings for {dataset_name}...")
+    
+    # ✅ Extract text based on dataset structure
+    if dataset_name == "sales":
+        texts = [" ".join(row.values()) for row in dataset["train"]]
+    elif dataset_name == "blended":
+        texts = [" ".join(row["free_messages"] + row["guided_messages"]) for row in dataset["train"]]
+    elif dataset_name == "dialog":
+        texts = [" ".join(row["dialog"]) for row in dataset["train"]]
+    elif dataset_name == "multiwoz":
+        texts = [" ".join(row["turns"]["utterance"]) for row in dataset["train"]]
+    else:
+        log(f"⚠️ Unknown dataset structure for {dataset_name}!")
+        texts = []
 
-    try:
-        index = faiss.read_index(index_path)
-        num_vectors = index.ntotal
-        dim = index.d
-        return f"📊 FAISS index contains {num_vectors} vectors.\n✅ Embedding dimension: {dim}"
-    except Exception as e:
-        return f"❌ ERROR: Failed to load FAISS index - {e}"
+    log(f"✅ Extracted {len(texts)} texts from {dataset_name}.")
 
-log("🔍 Checking FAISS embeddings...")
-faiss_status = check_faiss()
-log(faiss_status)
+    # ✅ 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)
 
-# ✅ Step 3: Initialize Chatbot
-client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+        # ✅ Log progress
+        log(f"🚀 Processed {i + len(batch)}/{len(texts)} embeddings for {dataset_name}...")
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}]
+        # ✅ Simulate delay for monitoring
+        time.sleep(1)
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    # ✅ Convert list of numpy arrays to a single numpy array
+    all_embeddings = np.vstack(all_embeddings)
+    return all_embeddings
 
-    messages.append({"role": "user", "content": message})
-    response = ""
+# ✅ Save embeddings to FAISS with unique filename
+def save_embeddings_to_faiss(embeddings, index_name="my_embeddings"):
+    index_file = f"{index_name}.faiss"
 
-    for message in client.chat_completions(
-        messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
-    ):
-        token = message["choices"][0]["delta"]["content"]
-        response += token
-        yield response
+    # ✅ Create new FAISS index
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(np.array(embeddings).astype(np.float32))
 
-# ✅ Step 4: Start Chatbot Interface
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
-)
+    # ✅ Save FAISS index
+    faiss.write_index(index, index_file)
+    log(f"✅ Saved FAISS index: {index_file}")
 
-log("✅ All systems go! Launching chatbot...")
-if __name__ == "__main__":
-    demo.launch()  
+# ✅ Run embedding process for all datasets
+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.")