NimrodDev commited on
Commit
ea6472f
·
1 Parent(s): e8c96fe

use optimum ONNX MiniLM via low-level wrapper (no disk)

Browse files
Files changed (1) hide show
  1. rag.py +15 -6
rag.py CHANGED
@@ -103,18 +103,27 @@ def get_texts() -> List[str]:
103
  print(f"⚠ Dataset fetch failed: {e} – using empty corpus")
104
  return []
105
 
106
- ## ------------------------------------------------------------------------------rtutu
107
- # ------------------------------------------------------------------
108
  # ------------------------------------------------------------------
109
  @lru_cache(maxsize=1)
110
  def get_vectorstore() -> FAISS:
111
  texts = get_texts()
112
 
113
  # --- FINAL: use optimum ONNX MiniLM (already on disk) -----------------
114
- from langchain_community.embeddings import OptimumSentenceEmbeddings
115
- embeddings = OptimumSentenceEmbeddings(
116
- model_name="optimum/all-MiniLM-L6-v2"
117
- )
 
 
 
 
 
 
 
 
 
 
 
118
  # ------------------------------------------------------------------------
119
 
120
  if not texts: # no data → empty FAISS
 
103
  print(f"⚠ Dataset fetch failed: {e} – using empty corpus")
104
  return []
105
 
 
 
106
  # ------------------------------------------------------------------
107
  @lru_cache(maxsize=1)
108
  def get_vectorstore() -> FAISS:
109
  texts = get_texts()
110
 
111
  # --- FINAL: use optimum ONNX MiniLM (already on disk) -----------------
112
+ import os
113
+ from optimum.pipelines import pipeline
114
+ from langchain.embeddings.base import Embeddings
115
+
116
+ class OptimumMiniLM(Embeddings):
117
+ def __init__(self):
118
+ self.pipe = pipeline("feature-extraction",
119
+ model="optimum/all-MiniLM-L6-v2",
120
+ device="cpu")
121
+ def embed_documents(self, texts):
122
+ return [self.pipe(t)[0][0] for t in texts]
123
+ def embed_query(self, text):
124
+ return self.embed_documents([text])[0]
125
+
126
+ embeddings = OptimumMiniLM()
127
  # ------------------------------------------------------------------------
128
 
129
  if not texts: # no data → empty FAISS