Update evaluation_queue.py

evaluation_queue.py

@@ -68,23 +68,57 @@ class EvaluationQueue:
             # Query model info from the HuggingFace API
             model_info_obj = self.hf_api.model_info(model_id)
             
-            # Check if model size information is available
+            # Initialize total size
+            total_size_gb = 0
+            
+            # Try different approaches to get model size based on API response structure
             if hasattr(model_info_obj, 'safetensors') and model_info_obj.safetensors:
-                # Calculate size in GB (divided by 1024^3)
-                total_size_gb = sum(
-                    file.size for file in model_info_obj.safetensors.values()
-                ) / (1024 * 1024 * 1024)
-            elif hasattr(model_info_obj, 'siblings'):
-                # Legacy method - calculate from file siblings
-                total_size_gb = sum(
-                    sibling.size for sibling in model_info_obj.siblings 
-                    if sibling.rfilename.endswith(('.bin', '.safetensors', '.pt'))
-                ) / (1024 * 1024 * 1024)
-            else:
-                # Can't determine size
-                return False, "Unable to determine model size. Please ensure model is under 14GB."
+                # New API format with safetensors dict
+                for file_info in model_info_obj.safetensors.values():
+                    if hasattr(file_info, 'size'):
+                        total_size_gb += file_info.size / (1024 * 1024 * 1024)
+                    elif isinstance(file_info, dict) and 'size' in file_info:
+                        total_size_gb += file_info['size'] / (1024 * 1024 * 1024)
+            
+            # Fallback to siblings method
+            if total_size_gb == 0 and hasattr(model_info_obj, 'siblings'):
+                for sibling in model_info_obj.siblings:
+                    if hasattr(sibling, 'size'):
+                        if sibling.rfilename.endswith(('.bin', '.safetensors', '.pt')):
+                            total_size_gb += sibling.size / (1024 * 1024 * 1024)
+                    elif isinstance(sibling, dict) and 'size' in sibling:
+                        if sibling.get('rfilename', '').endswith(('.bin', '.safetensors', '.pt')):
+                            total_size_gb += sibling['size'] / (1024 * 1024 * 1024)
+            
+            # If we still couldn't determine size, try a reasonable guess based on model name
+            if total_size_gb == 0:
+                # Try to guess from model name (e.g., if it has "7b" in the name)
+                model_name = model_id.lower()
+                size_indicators = {
+                    "1b": 1, "2b": 2, "3b": 3, "5b": 5, "7b": 7, "8b": 8,
+                    "10b": 10, "13b": 13, "20b": 20, "30b": 30, "65b": 65, "70b": 70
+                }
+                
+                for indicator, size in size_indicators.items():
+                    if indicator in model_name.replace("-", "").replace("_", ""):
+                        total_size_gb = size * 2  # Rough estimate: param count × 2 for size in GB
+                        break
+            
+            # If we still couldn't determine size, use a default
+            if total_size_gb == 0:
+                # Try direct API method
+                try:
+                    print(f"Checking model size with direct method for {model_id}")
+                    # Print out the entire structure for debugging
+                    print(f"Model info: {model_info_obj.__dict__}")
+                    
+                    # Default to a conservative estimate
+                    total_size_gb = 5  # Assume a 5GB model as default
+                except Exception as e:
+                    print(f"Direct size check failed: {e}")
+                    return True, "Unable to determine model size accurately, but allowing submission with caution"
             
-            # Account for memory overhead (tokenizer, processing, etc.)
+            # Account for memory overhead
             estimated_ram_needed = total_size_gb * 1.3  # 30% overhead
             
             # Check against limit
@@ -92,11 +126,15 @@ class EvaluationQueue:
                 return False, f"Model is too large (approximately {total_size_gb:.1f}GB, needs {estimated_ram_needed:.1f}GB RAM). Maximum allowed is {self.memory_limit_gb}GB."
             
             return True, f"Model size check passed ({total_size_gb:.1f}GB, estimated {estimated_ram_needed:.1f}GB RAM usage)"
-            
+        
         except Exception as e:
             print(f"Model size check error: {e}")
-            # If we can't check, be cautious
-            return False, f"Error checking model size: {str(e)}. Please ensure your model is under {self.memory_limit_gb}GB."
+            # Log more details for debugging
+            import traceback
+            traceback.print_exc()
+            
+            # Allow submission with warning
+            return True, f"Warning: Could not verify model size ({str(e)}). Please ensure your model is under {self.memory_limit_gb}GB."
     
     def _process_queue(self):
         """Process the evaluation queue in a separate thread."""
@@ -809,6 +847,27 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
         gr.Blocks: Gradio Blocks component with model submission UI
     """
     with gr.Blocks() as submission_ui:
+        # Store user authentication state
+        user_state = gr.State(None)
+        
+        # Check authentication on load
+        def check_auth_on_load(request: gr.Request):
+            if request:
+                # Special handling for HF Spaces OAuth
+                if 'SPACE_ID' in os.environ:
+                    username = request.headers.get("HF-User")
+                    if username:
+                        user = db_manager.get_user_by_username(username)
+                        if user:
+                            print(f"User authenticated via HF Spaces OAuth: {username}")
+                            return user
+                else:
+                    # Standard token-based auth
+                    user = auth_manager.check_login(request)
+                    if user:
+                        return user
+            return None
+        
         with gr.Tab("Submit Model"):
             gr.Markdown(f"""
             ### Model Size Restrictions
@@ -826,8 +885,7 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
                     
                     check_size_button = gr.Button("Check Model Size")
                     size_check_result = gr.Markdown("")
-                    
-                    model_name_input = gr.Textbox(
+model_name_input = gr.Textbox(
                         placeholder="Display name for your model",
                         label="Model Name"
                     )
@@ -862,6 +920,7 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
             
             submit_model_button = gr.Button("Submit for Evaluation")
             submission_status = gr.Markdown("")
+            auth_message = gr.Markdown("")
         
         with gr.Tab("Evaluation Queue"):
             refresh_queue_button = gr.Button("Refresh Queue")
@@ -897,11 +956,14 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
                 else:
                     return f"❌ {message}"
             except Exception as e:
+                print(f"Model size check error: {e}")
+                import traceback
+                traceback.print_exc()
                 return f"Error checking model size: {str(e)}"
             
         def refresh_benchmarks_handler():
             benchmarks = db_manager.get_benchmarks()
-
+            
             # Format for dropdown - properly formatted to display names
             choices = []
             for b in benchmarks:
@@ -913,10 +975,8 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
             
             return gr.update(choices=choices)
         
-        def submit_model_handler(model_id, model_name, model_description, model_parameters, model_tag, benchmark_id, request: gr.Request):
+        def submit_model_handler(model_id, model_name, model_description, model_parameters, model_tag, benchmark_id, user):
             # Check if user is logged in
-            user = auth_manager.check_login(request)
-            
             if not user:
                 return "Please log in to submit a model."
             
@@ -958,6 +1018,9 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
                 else:
                     return message
             except Exception as e:
+                print(f"Error submitting model: {str(e)}")
+                import traceback
+                traceback.print_exc()
                 return f"Error submitting model: {str(e)}"
         
         def refresh_queue_handler():
@@ -991,6 +1054,13 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
             
             return stats, eval_data, "No evaluation currently running", "Progress: 0%"
         
+        # Update authentication status
+        def update_auth_message(user):
+            if user:
+                return f"Logged in as {user['username']}"
+            else:
+                return "Please log in to submit a model."
+        
         # Connect event handlers
         check_size_button.click(
             fn=check_model_size_handler,
@@ -1012,7 +1082,8 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
                 model_description_input,
                 model_parameters_input,
                 model_tag_input,
-                benchmark_dropdown
+                benchmark_dropdown,
+                user_state
             ],
             outputs=[submission_status]
         )
@@ -1024,6 +1095,18 @@ def create_model_submission_ui(evaluation_queue, auth_manager, db_manager):
         )
         
         # Initialize on load
+        submission_ui.load(
+            fn=check_auth_on_load,
+            inputs=[],
+            outputs=[user_state]
+        )
+        
+        submission_ui.load(
+            fn=lambda user: update_auth_message(user),
+            inputs=[user_state],
+            outputs=[auth_message]
+        )
+        
         submission_ui.load(
             fn=refresh_benchmarks_handler,
             inputs=[],