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

# Initialize the FastAPI app
app = FastAPI()

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 my dashboard"}

# Define the API endpoint
@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, decimals=2).tolist()

        # Get the number of dimensions
        dimensions = len(rounded_embedding)

        # Return structured response
        return {
            "dimensions": dimensions,
            "embeddings": [rounded_embedding]  # Wrap the embedding inside a list
        }
    except Exception as e:
        # Handle any errors
        raise HTTPException(status_code=500, detail=str(e))

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