File size: 1,667 Bytes
7430631 8d520f8 7430631 3ee1de0 b7eccce a260b1b b7eccce be4dd45 7430631 |
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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, root_validator
from transformers import AutoModel
from typing import List
import os
if platform.system() == "Windows":
print("Windows detected. Assigning cache directory to Transformers in AppData\Local.")
transformers_cache_directory = os.path.join(os.getenv('LOCALAPPDATA'), 'transformers_cache')
if not os.path.exists(transformers_cache_directory):
try:
os.mkdir(transformers_cache_directory)
print(f"First launch. Directory '{transformers_cache_directory}' created successfully.")
except OSError as e:
print(f"Error creating directory '{transformers_cache_directory}': {e}")
else:
print(f"Directory '{transformers_cache_directory}' already exists.")
os.environ['TRANSFORMERS_CACHE'] = transformers_cache_directory
print("Environment variable assigned.")
del transformers_cache_directory
else:
print("Windows not detected. Assignment of Transformers cache directory not necessary.")
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',token = "hf_GkUomApayMBJteRvrjvslfyLRvfp QRckba".replace(" ", ""), trust_remote_code=True)
app = FastAPI()
class Validation(BaseModel):
prompt: List[str]
#Endpoint
@app.post("/jina_embedding")
async def chaatie_agent(item: Validation):
# Assuming model.encode returns a list of numpy arrays (one for each prompt)
embeddings = model.encode(item.prompt)
# Convert each numpy array in the list to a list
embeddings_list = [embedding.tolist() for embedding in embeddings]
return {"embeddings": embeddings_list}
|