Spaces:
Running
Running
File size: 759 Bytes
acedc04 |
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 |
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from typing import List
# Initialize the model
model = SentenceTransformer("PartAI/Tooka-SBERT")
# Create the FastAPI app
app = FastAPI()
# Pydantic model for input data
class TextInput(BaseModel):
sentences: List[str]
@app.get('/')
def index():
return {'message': 'Sentence embedding API.'}
# Endpoint to get embeddings
@app.post("/get_embeddings")
async def get_embeddings(input_data: TextInput):
# Get embeddings for the input sentences
embeddings = model.encode(input_data.sentences)
return {"embeddings": embeddings.tolist()}
# To run the app, save this code to a file, and then run `uvicorn filename:app --reload`
|