molehh commited on
Commit
9743eb6
·
1 Parent(s): 8e8b80a

first commit

Browse files
Files changed (3) hide show
  1. Dockerfile +22 -0
  2. main.py +48 -0
  3. requirements.txt +6 -0
Dockerfile ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9
2
+
3
+ WORKDIR /app
4
+ COPY . /app
5
+
6
+ ENV HF_HOME=/app/.cache
7
+
8
+ RUN mkdir -p /app/.cache/huggingface/hub && \
9
+ chmod -R 777 /app/.cache && \
10
+ chmod -R 777 /app/.cache/huggingface
11
+
12
+
13
+
14
+ RUN pip install --upgrade pip
15
+ RUN pip install --no-cache-dir -r requirements.txt
16
+
17
+ COPY --chown=user ./requirements.txt requirements.txt
18
+ RUN pip install --no-cache-dir --upgrade -r requirements.txt
19
+
20
+ EXPOSE 7860
21
+
22
+ CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
main.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
10
+ model = SentenceTransformer("Alibaba-NLP/gte-base-en-v1.5", trust_remote_code=True)
11
+
12
+ # Define the request body schema
13
+ class TextInput(BaseModel):
14
+ text: str
15
+
16
+ # Home route
17
+ @app.get("/")
18
+ async def home():
19
+ return {"message": "Welcome to the embedding SMS API. Use /docs to post SMS text and get dimensions."}
20
+
21
+ # Define the API endpoint for generating embeddings
22
+ @app.post("/embed")
23
+ async def generate_embedding(text_input: TextInput):
24
+ """
25
+ Generate a 768-dimensional embedding for the input text.
26
+ Returns the embedding in a structured format with rounded values.
27
+ """
28
+ try:
29
+ # Generate the embedding
30
+ embedding = model.encode(text_input.text, convert_to_tensor=True).cpu().numpy()
31
+
32
+ # Round embedding values to 2 decimal places
33
+ rounded_embedding = np.round(embedding, 2).tolist()
34
+
35
+ # Return structured response
36
+ return {
37
+ "dimensions": len(rounded_embedding),
38
+ "embeddings": [rounded_embedding]
39
+ }
40
+
41
+ except Exception as e:
42
+ # Handle any errors
43
+ raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
44
+
45
+ # Run the FastAPI app
46
+ if __name__ == "__main__":
47
+ import uvicorn
48
+ uvicorn.run(app, host="0.0.0.0", port=7860)
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ fastapi
2
+ pandas
3
+ scikit-learn
4
+ joblib
5
+ uvicorn
6
+ sentence-transformers