main.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import numpy as np

# Initialize the FastAPI app
app = FastAPI()

# Load the pre-trained SentenceTransformer model
model = SentenceTransformer("Alibaba-NLP/gte-base-en-v1.5", trust_remote_code=True)

# Define the request body schema
class TextInput(BaseModel):
    text: str

# Home route
@app.get("/")
async def home():
    return {"message": "welcome to home page"}

# Define the API endpoint for generating embeddings
@app.post("/embed")
async def generate_embedding(text_input: TextInput):
    """
    Generate a 768-dimensional embedding for the input text.
    Returns the embedding in a structured format with rounded values.
    """
    try:
        # Generate the embedding
        embedding = model.encode(text_input.text, convert_to_tensor=True).cpu().numpy()

        # Round embedding values to 2 decimal places
        rounded_embedding = np.round(embedding, 2).tolist()

        # Return structured response
        return {
            "dimensions": len(rounded_embedding),
            "embeddings": [rounded_embedding]
        }

    except Exception as e:
        # Handle any errors
        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")

# Run the FastAPI app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)