speed up inference

app.py

@@ -47,11 +47,11 @@ print("Device name:", torch.cuda.get_device_name(0))
 
 # If lower memory usage needed: 
 
-# bnb_config = BitsAndBytesConfig(
-#     load_in_4bit=True,
-#     bnb_4bit_use_double_quant=True,
-#     bnb_4bit_quant_type="nf4"
-# )
+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,
@@ -81,8 +81,8 @@ def load_model():
         local_files_only=True,
         torch_dtype=dtype,
         trust_remote_code=True,
-        # quantization_config=bnb_config,
-        # no_split_module_classes=["Block"],
+        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)
@@ -94,6 +94,9 @@ def load_model():
     )
     print("✅ Model and tokenizer loaded from", MODEL_CACHE)
 
+if torch.__version__ >= "2.0":
+    model = torch.compile(model)
+
 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()
@@ -103,6 +106,7 @@ def predict_NuExtract(texts, template, batch_size=1, max_length=5096, max_new_to
         "### Instruction:\n"
         "Remplis la template JSON avec les informations extraits du texte.\n"
         "Exemples types de formations :  CAP Boucherie, Licence Pro Métiers de l’Énergétique, Baccalauréat Général\n"
+        "Si le text contient des mentions de diplomes ou de titres de formations, certificats ou Attestation on les considère comme un item éducation\n"
         "Exemples catégories de formations :  Transport, énergie, langues, esthétique\n"
         "Exemples mobilités :  permis B, permis C, permis D. si y'a juste la mention de permis on considère que c'est le permis B\n"
         "Output *only* the completed JSON.\n"
@@ -126,7 +130,7 @@ def predict_NuExtract(texts, template, batch_size=1, max_length=5096, max_new_to
                 max_length=max_length
             ).to(device)
             print(f"Generating outputs with model for batch {i//batch_size+1}...", flush=True)
-            ids = model.generate(**enc, max_new_tokens=max_new_tokens, use_cache=False)
+            ids = model.generate(**enc, max_new_tokens=max_new_tokens, num_beams=1, use_cache=True)
             outputs += tokenizer.batch_decode(ids, skip_special_tokens=True)
     print("Outputs generated.", flush=True)
     elapsed = time.perf_counter() - start_time