Spaces:
Paused
Paused
File size: 1,252 Bytes
91bee69 d92c861 3a02c2a 07fb065 572cc27 d92c861 66e97f3 07fb065 aaf0100 07fb065 aaf0100 91bee69 3a02c2a |
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 48 49 50 |
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoModel, AutoTokenizer
import torch
device = torch.device("cpu")
# Load the model and tokenizer
model = AutoModel.from_pretrained(
"nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
"nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True
)
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.post("/get_embeding")
async def get_embeding(chunk):
# Tokenize the input text
inputs = tokenizer(chunk, return_tensors="pt")
# Generate embeddings
with torch.no_grad():
outputs = model(**inputs)
# The embeddings can be found in the 'last_hidden_state'
embeddings = outputs.last_hidden_state
# Optionally, you can average the token embeddings to get a single vector for the sentence
sentence_embedding = torch.mean(embeddings, dim=1)
#print(sentence_embedding)
return sentence_embedding.tolist()
|