enforce loading models locally, so no time is wasted downloading them (…

…#4)

Dockerfile

@@ -19,8 +19,7 @@ RUN python -c 'from optimum.onnxruntime import ORTModelForSequenceClassification
     model = ORTModelForSequenceClassification.from_pretrained("dcferreira/detoxify-optimized")'
 RUN python -c 'from transformers import pipeline;\
     model_path = "cardiffnlp/twitter-xlm-roberta-base-sentiment";\
-    sentiment_task = pipeline("sentiment-analysis", model=model_path, tokenizer=model_path);\
-    print(sentiment_task("love you!"))'
+    sentiment_task = pipeline("sentiment-analysis", model=model_path, tokenizer=model_path)'
 
 # copy the serving code and our database
 COPY app.py /code

app.py

@@ -2,6 +2,7 @@
 
 from dotenv import load_dotenv
 from fastapi import FastAPI
+from huggingface_hub import scan_cache_dir
 from loguru import logger
 from pydantic import HttpUrl
 from pymongo import MongoClient
@@ -21,6 +22,9 @@
 )
 db: Database = client[os.environ["DATABASE_NAME"]]
 
+if os.environ.get("DEBUG", False) == "1":
+    print(scan_cache_dir())
+
 
 @app.post("/fetch")
 async def fetch(url: HttpUrl) -> int:

docker-compose.yml

@@ -11,5 +11,10 @@ services:
     environment:
       - MONGODB_URL=mongodb://mongo:27017
       - DATABASE_NAME=test_database
+      - DEBUG=1
     ports:
       - "5000:8080"
+    deploy:
+      resources:
+        limits:
+          memory: 2G