Update ingest.py

ingest.py

@@ -1,13 +1,14 @@
 # ingest.py
 """
-Create / rebuild FAISS vector stores for Czech and English PDFs.
+Create FAISS indices for Czech and English PDFs.
 
-Default behaviour (matches main.py):
-  • English embeddings : sentence-transformers/all-MiniLM-L6-v2   (384-d)
-  • Czech   embeddings : Seznam/retromae-small-cs                 (768-d)
+  Default (matches backend/main.py):
+      • English embeddings : sentence-transformers/all-MiniLM-L6-v2  (384-d)
+      • Czech   embeddings : Seznam/retromae-small-cs               (768-d)
 
-Set use_openai=True if you really want to produce an English store
-with OpenAI's 3 072-d 'text-embedding-3-large' vectors.
+  If you still need a legacy English store with OpenAI
+  `text-embedding-3-large` (3 072-d), instantiate with
+      use_openai_embeddings=True   and pass OPENAI_API_KEY.
 """
 
 from pathlib import Path
@@ -16,10 +17,10 @@ from typing import List
 from langchain_community.vectorstores import FAISS
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.document_loaders import DirectoryLoader, PyPDFLoader
-from langchain.embeddings import (
-    OpenAIEmbeddings,
-    HuggingFaceEmbeddings,
-)
+
+# ← updated import (fixes deprecation warning) ----------------------[2][3]
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+from langchain.embeddings import OpenAIEmbeddings
 
 
 class Ingest:
@@ -27,26 +28,28 @@ class Ingest:
     def __init__(
         self,
         *,
-        # --- embeddings ----------------------------------------------------
-        english_hf_model: str = "sentence-transformers/all-MiniLM-L6-v2",
-        czech_hf_model: str   = "Seznam/retromae-small-cs",
-        english_oa_model: str = "text-embedding-3-large",
-        use_openai: bool      = False,              # flip to keep legacy store
-        openai_api_key: str | None = None,
-        # --- chunking ------------------------------------------------------
+        # names must stay exactly like in backend/main.py
+        english_embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
+        czech_embedding_model:   str = "Seznam/retromae-small-cs",
+        # optional OpenAI path
+        use_openai_embeddings:   bool = False,
+        openai_api_key:          str | None = None,
+        openai_embedding_model:  str = "text-embedding-3-large",
+        # chunking
         chunk: int = 512,
         overlap: int = 256,
-        # --- paths ---------------------------------------------------------
+        # folders
         english_store: str = "stores/english_512",
         czech_store:   str = "stores/czech_512",
         data_english:  str = "data/english",
         data_czech:    str = "data/czech",
     ):
-        self.use_openai   = use_openai
-        self.oa_key       = openai_api_key
-        self.english_hf   = english_hf_model
-        self.czech_hf     = czech_hf_model
-        self.english_oa   = english_oa_model
+        self.english_embedding_model = english_embedding_model
+        self.czech_embedding_model   = czech_embedding_model
+
+        self.use_openai_embeddings   = use_openai_embeddings
+        self.openai_api_key          = openai_api_key
+        self.openai_embedding_model  = openai_embedding_model
 
         self.chunk   = chunk
         self.overlap = overlap
@@ -58,89 +61,84 @@ class Ingest:
 
     # --------------------------- helpers ---------------------------------- #
     @staticmethod
-    def _loader(folder: Path):
+    def _load(folder: Path):
         return DirectoryLoader(
             str(folder),
             recursive=True,
-            show_progress=True,
             loader_cls=PyPDFLoader,
+            show_progress=True,
             use_multithreading=True,
         ).load()
 
     @staticmethod
     def _split(docs: List, chunk: int, overlap: int):
-        splitter = RecursiveCharacterTextSplitter(chunk_size=chunk,
-                                                  chunk_overlap=overlap)
-        return splitter.split_documents(docs)
+        return RecursiveCharacterTextSplitter(
+            chunk_size=chunk, chunk_overlap=overlap
+        ).split_documents(docs)
 
     # --------------------------- English ---------------------------------- #
     def ingest_english(self):
-        if self.use_openai:
-            if not self.oa_key:
-                raise ValueError("OpenAI API key is required for OpenAI embeddings.")
-            embedding = OpenAIEmbeddings(
-                openai_api_key=self.oa_key,
-                model=self.english_oa,
+        if self.use_openai_embeddings:
+            if not self.openai_api_key:
+                raise ValueError("OPENAI_API_KEY missing for OpenAI embeddings.")
+            embed = OpenAIEmbeddings(
+                openai_api_key=self.openai_api_key,
+                model=self.openai_embedding_model,
             )
-            mode = f"OpenAI ({self.english_oa}) 3072-d"
+            mode = f"OpenAI ({self.openai_embedding_model}) 3 072-d"
         else:
-            embedding = HuggingFaceEmbeddings(
-                model_name=self.english_hf,
+            embed = HuggingFaceEmbeddings(
+                model_name=self.english_embedding_model,
                 model_kwargs={"device": "cpu"},
                 encode_kwargs={"normalize_embeddings": False},
             )
-            mode = f"HuggingFace ({self.english_hf}) " \
-                   f"{embedding.client.get_sentence_embedding_dimension()}-d"
+            dim  = embed.client.get_sentence_embedding_dimension()
+            mode = f"HuggingFace ({self.english_embedding_model}) {dim}-d"
 
-        print(f"\n─ Ingest EN: {mode}")
-        docs  = self._loader(self.data_english)
-        texts = self._split(docs, self.chunk, self.overlap)
-
-        db = FAISS.from_documents(texts, embedding)
-        db.save_local(str(self.english_store))
-        print("✓ English store written to", self.english_store, "\n")
+        print(f"\n── Building English index with {mode}")
+        texts = self._split(self._load(self.data_english), self.chunk, self.overlap)
+        FAISS.from_documents(texts, embed).save_local(str(self.english_store))
+        print("✓ English store saved to", self.english_store, "\n")
 
     # --------------------------- Czech ------------------------------------ #
     def ingest_czech(self):
-        embedding = HuggingFaceEmbeddings(
-            model_name=self.czech_hf,
+        embed = HuggingFaceEmbeddings(
+            model_name=self.czech_embedding_model,
             model_kwargs={"device": "cpu"},
             encode_kwargs={"normalize_embeddings": False},
         )
-        dim = embedding.client.get_sentence_embedding_dimension()
-        print(f"\n─ Ingest CZ: HuggingFace ({self.czech_hf}) {dim}-d")
-
-        docs  = self._loader(self.data_czech)
-        texts = self._split(docs, self.chunk, self.overlap)
+        dim = embed.client.get_sentence_embedding_dimension()
+        print(f"\n── Building Czech index with HuggingFace "
+              f"({self.czech_embedding_model}) {dim}-d")
+        texts = self._split(self._load(self.data_czech), self.chunk, self.overlap)
+        FAISS.from_documents(texts, embed).save_local(str(self.czech_store))
+        print("✓ Czech store saved to", self.czech_store, "\n")
 
-        db = FAISS.from_documents(texts, embedding)
-        db.save_local(str(self.czech_store))
-        print("✓ Czech store written to", self.czech_store, "\n")
 
-
-# -------------------- quick CLI helper ------------------------------------ #
+# ───────────── CLI helper (optional) ───────────── #
 if __name__ == "__main__":
     """
-    Examples:
-        # build both stores with default HF encoders (no OpenAI)
-        python ingest.py
-
-        # build English store with OpenAI encoder (keeps 3 072-d index)
-        OPENAI_API_KEY=sk-... python ingest.py --openai
+    Examples
+    --------
+    python ingest.py                 # builds both stores (OSS embeddings)
+    OPENAI_API_KEY=sk-... \
+    python ingest.py --openai en     # rebuild English with OpenAI encoder
     """
     import argparse, os
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--openai", action="store_true",
-                        help="Use OpenAI embeddings for English.")
-    parser.add_argument("--only", choices=["en", "cz"],
-                        help="Ingest only that language.")
-    args = parser.parse_args()
+    p = argparse.ArgumentParser()
+    p.add_argument("--openai", action="store_true",
+                   help="Use OpenAI embeddings for English store.")
+    p.add_argument("lang", nargs="?", choices=["en", "cz"],
+                   help="Only ingest this language.")
+    args = p.parse_args()
 
-    ing = Ingest(use_openai=args.openai,
-                 openai_api_key=os.getenv("OPENAI_API_KEY"))
+    ing = Ingest(
+        use_openai_embeddings=args.openai,
+        openai_api_key=os.getenv("OPENAI_API_KEY"),
+    )
 
-    if args.only in (None, "en"):
+    if args.lang in (None, "en"):
         ing.ingest_english()
-    if args.only in (None, "cz"):
+    if args.lang in (None, "cz"):
         ing.ingest_czech()