moraxgiga commited on
Commit
be4dd45
1 Parent(s): 7430631

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +19 -0
main.py CHANGED
@@ -3,5 +3,24 @@ from pydantic import BaseModel, root_validator
3
  from transformers import AutoModel
4
  from typing import List
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
 
7
 
 
3
  from transformers import AutoModel
4
  from typing import List
5
 
6
+ model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',token = access_token, trust_remote_code=True)
7
+
8
+ app = FastAPI()
9
+
10
+ class Validation(BaseModel):
11
+ prompt: List[str]
12
+
13
+
14
+ #Endpoint
15
+ @app.post("/jina_embedding")
16
+ async def chaatie_agent(item: Validation):
17
+ # Assuming model.encode returns a list of numpy arrays (one for each prompt)
18
+ embeddings = model.encode(item.prompt)
19
+
20
+ # Convert each numpy array in the list to a list
21
+ embeddings_list = [embedding.tolist() for embedding in embeddings]
22
+
23
+ return {"embeddings": embeddings_list}
24
+
25
 
26