size checker depends on model precision

src/submission/submit.py

@@ -2,6 +2,9 @@ import json
 import os
 from datetime import datetime, timezone
 
+from dataclasses import dataclass
+from transformers import AutoConfig
+
 from src.display.formatting import styled_error, styled_message, styled_warning
 from src.envs import (
     API,
@@ -23,6 +26,30 @@ from src.submission.check_validity import (
 REQUESTED_MODELS = None
 USERS_TO_SUBMISSION_DATES = None
 
+@dataclass
+class ModelSizeChecker:
+    model: str
+    precision: str
+    model_size_in_b: float
+
+    def get_precision_factor(self):
+        if self.precision in ["float16", "bfloat16"]:
+            return 1
+        elif self.precision == "8bit":
+            return 2
+        elif self.precision == "4bit":
+            return 4
+        elif self.precision == "GPTQ":
+            config = AutoConfig.from_pretrained(self.model)
+            num_bits = int(config.quantization_config["bits"])
+            bits_to_precision_factor = {2: 8, 3: 6, 4: 4, 8: 2}
+            return bits_to_precision_factor.get(num_bits, 1)
+        else:
+            raise Exception(f"Unknown precision {self.precision}.")
+
+    def can_evaluate(self):
+        precision_factor = self.get_precision_factor()
+        return self.model_size_in_b <= 140 * precision_factor
 
 def add_new_eval(
     model: str,
@@ -72,12 +99,20 @@ def add_new_eval(
 
     # Check model size early
     model_size = get_model_size(model_info=model_info, precision=precision)
-    if model_size > 100:
-        return styled_error(f"Sadly, models this big ({model_size}b parameters) cannot be evaluated automatically at the moment on our cluster")
-
-    # Check for duplicate submission
-    if f"{model}_{model_info.sha}_{precision}" in REQUESTED_MODELS:
-        return styled_warning("This model has been already submitted.")
+    
+    # First check: Absolute size limit for float16 and bfloat16
+    if precision in ["float16", "bfloat16"] and model_size > 100:
+        return styled_error(f"Sadly, models larger than 100B parameters cannot be submitted in {precision} precision at this time. "
+                            f"Your model size: {model_size:.2f}B parameters.")
+
+    # Second check: Precision-adjusted size limit
+    size_checker = ModelSizeChecker(model=model, precision=precision, model_size_in_b=model_size)
+    
+    if not size_checker.can_evaluate():
+        precision_factor = size_checker.get_precision_factor()
+        max_size = 140 * precision_factor
+        return styled_error(f"Sadly, models this big ({model_size:.2f}B parameters) cannot be evaluated automatically "
+                            f"at the moment on our cluster. The maximum size for {precision} precision is {max_size:.2f}B parameters.")
 
     architecture = "?"
     # Is the model on the hub?
@@ -134,13 +169,13 @@ def add_new_eval(
         f.write(json.dumps(eval_entry))
 
     print("Uploading eval file")
-    API.upload_file(
-        path_or_fileobj=out_path,
-        path_in_repo=out_path.split("eval-queue/")[1],
-        repo_id=QUEUE_REPO,
-        repo_type="dataset",
-        commit_message=f"Add {model} to eval queue",
-    )
+    # API.upload_file(
+    #     path_or_fileobj=out_path,
+    #     path_in_repo=out_path.split("eval-queue/")[1],
+    #     repo_id=QUEUE_REPO,
+    #     repo_type="dataset",
+    #     commit_message=f"Add {model} to eval queue",
+    # )
 
     # Remove the local file
     os.remove(out_path)