Update main.py
Browse files
main.py
CHANGED
@@ -2,28 +2,35 @@ from fastapi import FastAPI, HTTPException
|
|
2 |
from pydantic import BaseModel, root_validator
|
3 |
from transformers import AutoModel
|
4 |
from typing import List
|
5 |
-
import os, platform
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
|
10 |
|
|
|
|
|
11 |
app = FastAPI()
|
12 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
class Validation(BaseModel):
|
14 |
prompt: List[str]
|
15 |
|
16 |
-
|
17 |
-
#Endpoint
|
18 |
@app.post("/jina_embedding")
|
19 |
async def generate_embeddings(item: Validation):
|
20 |
-
start_time = time.time()
|
21 |
embeddings = model.encode(item.prompt).tolist()
|
22 |
-
end_time = time.time()
|
23 |
time_taken = end_time - start_time # Calculate the time taken
|
24 |
|
25 |
return {
|
26 |
"embeddings": embeddings,
|
27 |
"time_taken": f"{time_taken:.2f} seconds",
|
28 |
-
"Number_of_sentence_processed"
|
29 |
}
|
|
|
2 |
from pydantic import BaseModel, root_validator
|
3 |
from transformers import AutoModel
|
4 |
from typing import List
|
5 |
+
import os, platform, time
|
6 |
+
from fastapi.middleware.cors import CORSMiddleware # Import CORSMiddleware
|
|
|
|
|
|
|
7 |
|
8 |
+
# Your model loading and app initialization code
|
9 |
+
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
|
10 |
app = FastAPI()
|
11 |
|
12 |
+
# Set up CORS middleware configuration
|
13 |
+
app.add_middleware(
|
14 |
+
CORSMiddleware,
|
15 |
+
allow_origins=["*"], # Allows all origins
|
16 |
+
allow_credentials=True,
|
17 |
+
allow_methods=["*"], # Allows all methods
|
18 |
+
allow_headers=["*"], # Allows all headers
|
19 |
+
)
|
20 |
+
|
21 |
+
# Your Validation model and endpoint definitions remain the same
|
22 |
class Validation(BaseModel):
|
23 |
prompt: List[str]
|
24 |
|
|
|
|
|
25 |
@app.post("/jina_embedding")
|
26 |
async def generate_embeddings(item: Validation):
|
27 |
+
start_time = time.time()
|
28 |
embeddings = model.encode(item.prompt).tolist()
|
29 |
+
end_time = time.time()
|
30 |
time_taken = end_time - start_time # Calculate the time taken
|
31 |
|
32 |
return {
|
33 |
"embeddings": embeddings,
|
34 |
"time_taken": f"{time_taken:.2f} seconds",
|
35 |
+
"Number_of_sentence_processed": len(item.prompt)
|
36 |
}
|