script to assign labels to annotators

scripts/distribute_workload.py

@@ -0,0 +1,170 @@
+import sys
+import os
+
+# Add project root to Python path
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+if project_root not in sys.path:
+    sys.path.insert(0, project_root)
+
+import math
+import random # Added for password generation
+from sqlalchemy.sql import func
+
+from utils.database import get_db
+from data.models import TTSData
+from data.repository.annotator_repo import AnnotatorRepo
+from data.repository.annotation_interval_repo import AnnotationIntervalRepo
+from utils.logger import Logger
+
+log = Logger()
+
+# --- Configuration ---
+# List of annotator names to ensure exist and assign work to
+ANNOTATOR_NAMES = ["shahab", "amir", "mohsen", "mahya", "najmeh", "sepehr", "zahra", "moghim", "amin"]
+# DEFAULT_PASSWORD is no longer used for new users, random password will be generated.
+
+def generate_random_password():
+    """Generates a random 4-digit numerical password."""
+    return str(random.randint(1000, 9999))
+
+def distribute_workload():
+    log.info("Starting workload distribution script...")
+    processed_annotators_details = [] # Stores dicts: {'annotator_obj': obj, 'password_display': str, 'assigned_start': int, 'assigned_end': int}
+
+    try:
+        with get_db() as db:
+            annot_repo = AnnotatorRepo(db)
+            interval_repo = AnnotationIntervalRepo(db)
+
+            # 1. Ensure all annotators exist, create if not, and collect details
+            log.info("Processing annotators...")
+            for name in ANNOTATOR_NAMES:
+                annotator = annot_repo.get_annotator_by_name(name)
+                password_to_display = "(existing user)"
+                
+                if not annotator:
+                    try:
+                        new_password = generate_random_password()
+                        log.info(f"Annotator '{name}' not found, creating with new password...")
+                        annotator = annot_repo.add_new_annotator(name, new_password)
+                        log.info(f"Annotator '{name}' (id={annotator.id}) created successfully with password '{new_password}'.")
+                        password_to_display = new_password
+                    except ValueError as e:
+                        log.warning(f"Could not create annotator '{name}' (likely already exists or other DB issue): {e}. Attempting to fetch again.")
+                        annotator = annot_repo.get_annotator_by_name(name) # Try fetching again
+                        if annotator:
+                            log.info(f"Found existing annotator '{name}' (id={annotator.id}) after creation attempt.")
+                        else:
+                            log.error(f"Failed to create or find annotator '{name}'. Skipping.")
+                            continue
+                else:
+                    log.info(f"Found existing annotator '{name}' (id={annotator.id}).")
+
+                if annotator:
+                    processed_annotators_details.append({
+                        'annotator_obj': annotator, 
+                        'password_display': password_to_display,
+                        'assigned_start': None,
+                        'assigned_end': None
+                    })
+            
+            if not processed_annotators_details:
+                log.error("No annotators processed or found. Exiting.")
+                return
+
+            # 2. Get total number of TTSData items
+            total_tts_items = db.query(func.count(TTSData.id)).scalar()
+            if total_tts_items is None or total_tts_items == 0:
+                log.info("No TTSData items found in the database. Nothing to assign.")
+                # Still print annotator info even if no items to assign
+                log.info("\\n--- Workload Distribution Summary ---")
+                for details in processed_annotators_details:
+                    log.info(f"Annotator: {details['annotator_obj'].name}, Assigned Range: N/A (No data items), Password: {details['password_display']}")
+                return
+            log.info(f"Total TTSData items found: {total_tts_items}")
+
+            # 3. Calculate distribution
+            num_annotators_for_assignment = len(processed_annotators_details)
+            if num_annotators_for_assignment == 0: # Should be caught by earlier check, but as a safeguard
+                log.error("No annotators available for assignment. Exiting.")
+                return
+
+            items_per_annotator_base = total_tts_items // num_annotators_for_assignment
+            remainder_items = total_tts_items % num_annotators_for_assignment
+
+            log.info(f"Distributing {total_tts_items} items among {num_annotators_for_assignment} annotators.")
+            log.info(f"Base items per annotator: {items_per_annotator_base}, Remainder: {remainder_items}")
+
+            # 4. Assign intervals
+            current_start_idx = 1 # Assuming TTSData IDs start from 1
+            for details_dict in processed_annotators_details:
+                annotator = details_dict['annotator_obj']
+                num_items_for_this_annotator = items_per_annotator_base
+                if remainder_items > 0:
+                    num_items_for_this_annotator += 1
+                    remainder_items -= 1
+                
+                if num_items_for_this_annotator == 0:
+                    log.info(f"Annotator '{annotator.name}' assigned 0 items (total items might be less than annotators or workload already distributed).")
+                    continue
+
+                current_end_idx = current_start_idx + num_items_for_this_annotator - 1
+
+                if current_end_idx > total_tts_items:
+                    current_end_idx = total_tts_items
+
+                if current_start_idx > current_end_idx:
+                    log.info(f"No items to assign to '{annotator.name}' (start_idx {current_start_idx} > end_idx {current_end_idx}).")
+                    continue
+
+                log.info(f"Attempting to assign interval [{current_start_idx}-{current_end_idx}] to '{annotator.name}' (id={annotator.id})")
+                try:
+                    existing_intervals = interval_repo.get_intervals_by_annotator(annotator.id)
+                    if existing_intervals:
+                        log.warning(f"Annotator '{annotator.name}' already has existing intervals. Skipping assignment to avoid conflicts. Manual review/cleanup of old intervals might be needed.")
+                        # current_start_idx = current_end_idx + 1 # This line should not be here if we skip the user for this round of assignment.
+                                                            # The items for this user won't be assigned and won't be passed to the next.
+                                                            # This means the total items might not be fully distributed if users are skipped.
+                                                            # For a full distribution even with skips, a more complex item re-allocation would be needed.
+                                                            # For now, skipped users mean their share is not re-distributed.
+                        continue # Skip this annotator for assignment
+
+                    assigned_interval = interval_repo.assign_interval_to_annotator(
+                        annotator_id=annotator.id,
+                        start_idx=current_start_idx,
+                        end_idx=current_end_idx,
+                        allow_overlap=False 
+                    )
+                    details_dict['assigned_start'] = assigned_interval.start_index
+                    details_dict['assigned_end'] = assigned_interval.end_index
+                    log.info(
+                        f"Successfully assigned interval [{details_dict['assigned_start']}-{details_dict['assigned_end']}] "
+                        f"to '{annotator.name}' (id={annotator.id})"
+                    )
+                except ValueError as e:
+                    log.error(f"Could not assign interval [{current_start_idx}-{current_end_idx}] to '{annotator.name}': {e}")
+                except Exception as e:
+                    log.error(f"An unexpected error occurred while assigning interval to '{annotator.name}': {e}")
+                
+                # Only advance current_start_idx if items were potentially assignable to *this* annotator
+                # If an annotator was skipped due to existing intervals, their share of items is not processed further in this loop.
+                current_start_idx = current_end_idx + 1 
+                if current_start_idx > total_tts_items:
+                    break 
+
+            # 5. Print summary
+            log.info("\\n--- Workload Distribution Summary ---")
+            for details in processed_annotators_details:
+                range_str = "N/A (assignment skipped or failed)"
+                if details['assigned_start'] is not None and details['assigned_end'] is not None:
+                    range_str = f"[{details['assigned_start']}-{details['assigned_end']}]"
+                
+                log.info(f"Annotator: {details['annotator_obj'].name}, Assigned Range: {range_str}, Password: {details['password_display']}")
+
+            log.info("Workload distribution script finished.")
+
+    except Exception as e:
+        log.error(f"An critical error occurred during workload distribution: {e}", exc_info=True)
+
+if __name__ == "__main__":
+    distribute_workload()