Speed Memory Usage optimizations

Dockerfile

@@ -3,7 +3,10 @@ FROM python:3.10-slim
 # 1) Variables HF avant tout
 ENV HF_HOME="/home/user/.cache/huggingface" \
     HF_HUB_CACHE="/home/user/.cache/huggingface/hub" \
-    TRANSFORMERS_CACHE="/home/user/.cache/huggingface/transformers"
+    TRANSFORMERS_CACHE="/home/user/.cache/huggingface/transformers" \
+    DOCLING_ARTIFACTS_PATH="/home/user/.cache/docling/models" \ 
+    OMP_NUM_THREADS=2
+
 
 # 2) Créer l’utilisateur non-root
 RUN useradd -m -u 1000 user

app.py

@@ -10,6 +10,8 @@ from dotenv import load_dotenv
 import tempfile
 from supabase import create_client
 from huggingface_hub import snapshot_download
+from transformers import BitsAndBytesConfig, AutoModelForCausalLM
+
 
 
 load_dotenv()
@@ -37,6 +39,23 @@ print(">>> MODEL CACHE PATH:", MODEL_CACHE, os.listdir(MODEL_CACHE))
 
 device = "gpu" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 dtype = torch.float16 if device in ("mps", "gpu") else torch.float32
+bnb_config = BitsAndBytesConfig(load_in_8bit=True)
+
+
+# If lower memory usage needed: 
+
+# bnb_config = BitsAndBytesConfig(
+#     load_in_4bit=True,
+#     bnb_4bit_use_double_quant=True,
+#     bnb_4bit_quant_type="nf4"
+# )
+# model = AutoModelForCausalLM.from_pretrained(
+#     MODEL_CACHE,
+#     quantization_config=bnb_config,
+#     device_map="auto",
+#     local_files_only=True,
+#     trust_remote_code=True
+# )
 
 @app.on_event("startup")
 def startup_supabase():
@@ -57,17 +76,21 @@ def load_model():
         MODEL_CACHE,
         local_files_only=True,
         torch_dtype=dtype,
-        trust_remote_code=True
+        trust_remote_code=True,
+        quantization_config=bnb_config,
+        no_split_module_classes=["Block"],
+        device_map="auto"
     ).to(device).eval()
     # tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_CACHE,
         local_files_only=True,
-        trust_remote_code=True
+        trust_remote_code=True,
+        device_map="auto"
     )
     print("✅ Model and tokenizer loaded from", MODEL_CACHE)
 
-def predict_NuExtract(texts, template, batch_size=10, max_length=5096, max_new_tokens=1024):
+def predict_NuExtract(texts, template, batch_size=1, max_length=5096, max_new_tokens=1024):
     print("Starting NuExtract prediction...", flush=True)
     start_time = time.perf_counter()
     template_str = json.dumps(json.loads(template), indent=4)