fix permission issue for cache, and add pooling

Dockerfile

@@ -2,11 +2,18 @@ FROM python:3.10-slim
 
 WORKDIR /app
 
+# Tạo thư mục cache riêng
+ENV TRANSFORMERS_CACHE=/app/cache
+ENV HF_HOME=/app/hf_home
+
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
 COPY app.py .
 
+# Tạo sẵn thư mục cache (tránh lỗi lần đầu chạy)
+RUN mkdir -p /app/cache && mkdir -p /app/hf_home
+
 EXPOSE 7860
 
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -1,27 +1,21 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModel
-import torch
+from sentence_transformers import SentenceTransformer
+import uvicorn
 
 app = FastAPI()
 
 # Load model
-model_name = "BAAI/bge-m3"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+model = SentenceTransformer("BAAI/bge-m3")
 
 class InputText(BaseModel):
     text: str
 
 @app.get("/")
-def root():
-    return {"message": "BAAI/bge-m3 embedding API is running."}
+def read_root():
+    return {"message": "BAAI/bge-m3 Sentence Embedding API is running."}
 
 @app.post("/embed")
 def get_embedding(data: InputText):
-    inputs = tokenizer(data.text, return_tensors="pt", padding=True, truncation=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        # Get CLS token or use pooling method
-        embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()
+    embedding = model.encode(data.text, normalize_embeddings=True).tolist()
     return {"embedding": embedding}

requirements.txt

@@ -1,4 +1,3 @@
-transformers==4.41.0
-torch
+sentence-transformers
 fastapi
 uvicorn