app.py

from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from typing import List

# Load the pre-trained sentence transformer model
bge_small_model = SentenceTransformer('BAAI/bge-small-en-v1.5', device="cpu")
all_mp_net_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device="cpu")

# Initialize FastAPI app
app = FastAPI()

# Request body model
class TextInput(BaseModel):
    text: List[str]  # List of sentences or text data
    model_name: str

# Route to calculate embeddings
@app.post("/get-embedding/")
async def get_embedding(input: TextInput):
    # Generate embeddings using the sentence transformer model
    if input.model_name == "BM":
        embeddings = all_mp_net_model.encode(input.text)
    else:
        embeddings = bge_small_model.encode(input.text)
    return {"embeddings": embeddings.tolist()}

@app.get("/status")
async def status():
    return {"status": "Server is up and running"}