File size: 1,628 Bytes
f3dc108 60c54da f3dc108 |
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 51 |
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 from Hugging Face
#model = SentenceTransformer("//huggingface.co/spaces/Kabila22/Kabilan_embedding_1", trust_remote_code=True)
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 embedded model"}
# 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) |