tharu22 commited on
Commit
f3dc108
·
1 Parent(s): a113d63
Files changed (1) hide show
  1. main.py +51 -0
main.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from pydantic import BaseModel
3
+ from sentence_transformers import SentenceTransformer
4
+ import numpy as np
5
+
6
+ # Initialize the FastAPI app
7
+ app = FastAPI()
8
+
9
+ # Load the pre-trained SentenceTransformer model from Hugging Face
10
+ #model = SentenceTransformer("//huggingface.co/spaces/Kabila22/Kabilan_embedding_1", trust_remote_code=True)
11
+ model = SentenceTransformer("Alibaba-NLP/gte-base-en-v1.5", trust_remote_code=True)
12
+
13
+ # Define the request body schema
14
+ class TextInput(BaseModel):
15
+ text: str
16
+
17
+ # Home route
18
+ @app.get("/")
19
+ async def home():
20
+ return {"message": "Welcome to embedding SMS API, use /docs to post SMS text and get dimensions"}
21
+
22
+ # Define the API endpoint
23
+ @app.post("/embed")
24
+ async def generate_embedding(text_input: TextInput):
25
+ """
26
+ Generate a 768-dimensional embedding for the input text.
27
+ Returns the embedding in a structured format with rounded values.
28
+ """
29
+ try:
30
+ # Generate the embedding
31
+ embedding = model.encode(text_input.text, convert_to_tensor=True).cpu().numpy()
32
+
33
+ # Round embedding values to 2 decimal places
34
+ rounded_embedding = np.round(embedding, decimals=2).tolist()
35
+
36
+ # Get the number of dimensions
37
+ dimensions = len(rounded_embedding)
38
+
39
+ # Return structured response
40
+ return {
41
+ "dimensions": dimensions,
42
+ "embeddings": [rounded_embedding] # Wrap the embedding inside a list
43
+ }
44
+ except Exception as e:
45
+ # Handle any errors
46
+ raise HTTPException(status_code=500, detail=str(e))
47
+
48
+ # Run the FastAPI app
49
+ if __name__ == "__main__":
50
+ import uvicorn
51
+ uvicorn.run(app, host="0.0.0.0", port=7860)