rohannsinghal commited on
Commit
1d35377
·
verified ·
1 Parent(s): 803fdb2

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +21 -0
handler.py ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from sentence_transformers import SentenceTransformer
2
+
3
+ class EndpointHandler:
4
+ def __init__(self, path=""):
5
+ # Here is the magic override that bypasses the Hugging Face bug
6
+ print("Initializing Nemotron 8B with trust_remote_code=True...")
7
+ self.model = SentenceTransformer("nvidia/llama-embed-nemotron-8b", trust_remote_code=True)
8
+ print("Model loaded successfully!")
9
+
10
+ def __call__(self, data):
11
+ """
12
+ This runs every time your Vectorize script sends text to the endpoint.
13
+ """
14
+ # Get the text from the API request
15
+ inputs = data.pop("inputs", data)
16
+
17
+ # Generate the math vectors
18
+ embeddings = self.model.encode(inputs)
19
+
20
+ # Return it as a standard Python list so your local script can read it
21
+ return embeddings.tolist()