AarnavNoble commited on
Commit
8923e30
·
verified ·
1 Parent(s): 2ae0744

Upload backend/ml/rag/embedder.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. backend/ml/rag/embedder.py +36 -0
backend/ml/rag/embedder.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Embeds text chunks using sentence-transformers (all-MiniLM-L6-v2).
3
+ Runs fully locally, no API key needed.
4
+ """
5
+
6
+ from sentence_transformers import SentenceTransformer
7
+ import numpy as np
8
+
9
+ MODEL_NAME = "all-MiniLM-L6-v2" # 384-dim, fast, good quality
10
+ _model = None
11
+
12
+
13
+ def get_model() -> SentenceTransformer:
14
+ global _model
15
+ if _model is None:
16
+ print(f"Loading embedding model: {MODEL_NAME}")
17
+ _model = SentenceTransformer(MODEL_NAME)
18
+ return _model
19
+
20
+
21
+ def embed_texts(texts: list[str]) -> np.ndarray:
22
+ """
23
+ Embed a list of strings. Returns (N, 384) float32 array.
24
+ """
25
+ model = get_model()
26
+ embeddings = model.encode(texts, batch_size=64, show_progress_bar=True, normalize_embeddings=True)
27
+ return embeddings.astype(np.float32)
28
+
29
+
30
+ def embed_query(query: str) -> np.ndarray:
31
+ """
32
+ Embed a single query string. Returns (384,) float32 array.
33
+ """
34
+ model = get_model()
35
+ embedding = model.encode([query], normalize_embeddings=True)
36
+ return embedding[0].astype(np.float32)