File size: 985 Bytes
d2aa2a5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
from sentence_transformers import SentenceTransformer
import numpy as np
import time

class Embedder:
    def __init__(self):
        print("Loading embedding model...")
        self.model = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2')
        print("Embedder loaded: paraphrase-MiniLM-L6-v")

    def embed(self, chunks, batch_size=32): 
        try:
            vectors = self.model.encode(
                chunks, 
                batch_size=batch_size, 
                show_progress_bar=False,
                convert_to_numpy=True
            )
            return vectors.tolist()
        except Exception as e:
            print(f"Embedding failed: {e}")
            raise e

    def embed_q(self, query):
        
        try:
            v = self.model.encode(query, convert_to_numpy=True)
            return v.tolist()
        except Exception as e:
            print(f"Query embedding failed: {e}")
            raise e