Refactor metric extraction and clean up unused code

Replaced calls to extract_consensus_metrics_from_discussion with direct use of _extract_metrics_from_text in consensus.py and removed the now-unused extract_consensus_metrics_from_discussion function. Cleaned up debug logging in providers/grok.py and removed unused metadata handling in utils.py. Fixed cache format count key from 'v1.0' to '1.0' in utils.py.

Author: cafferychen777

python/mllmcelltype/consensus.py

@@ -666,13 +666,11 @@ def process_controversial_clusters(
                     # Previously had consensus indicators check here, now using metrics extraction
 
                     # Extract consensus proportion and entropy values for the current round
-                    cp_value, h_value = extract_consensus_metrics_from_discussion(response)
+                    cp_value, h_value = _extract_metrics_from_text(response)[:2]
 
                     # If unable to extract from discussion, try to extract from consensus check response
                     if cp_value is None or h_value is None:
-                        cp_value, h_value = extract_consensus_metrics_from_discussion(
-                            consensus_response
-                        )
+                        cp_value, h_value = _extract_metrics_from_text(consensus_response)[:2]
 
                     # If still unable to extract, use default values
                     if cp_value is None:
@@ -710,7 +708,7 @@ def process_controversial_clusters(
                         )
 
                         # Extract CP and H from the discussion if available
-                        cp_value, h_value = extract_consensus_metrics_from_discussion(response)
+                        cp_value, h_value = _extract_metrics_from_text(response)[:2]
                         if cp_value is not None and h_value is not None:
                             updated_consensus_proportion[cluster_id] = cp_value
                             updated_entropy[cluster_id] = h_value
@@ -788,7 +786,7 @@ def process_controversial_clusters(
                 # if available
                 if rounds_history:
                     last_round = rounds_history[-1]
-                    cp_value, h_value = extract_consensus_metrics_from_discussion(last_round)
+                    cp_value, h_value = _extract_metrics_from_text(last_round)[:2]
                     if cp_value is not None and h_value is not None:
                         updated_consensus_proportion[cluster_id] = cp_value
                         updated_entropy[cluster_id] = h_value
@@ -806,7 +804,7 @@ def process_controversial_clusters(
                 # Extract metrics from the last round if available
                 if cluster_id not in updated_consensus_proportion and rounds_history:
                     last_round = rounds_history[-1]
-                    cp_value, h_value = extract_consensus_metrics_from_discussion(last_round)
+                    cp_value, h_value = _extract_metrics_from_text(last_round)[:2]
                     if cp_value is not None and h_value is not None:
                         updated_consensus_proportion[cluster_id] = cp_value
                         updated_entropy[cluster_id] = h_value
@@ -835,22 +833,6 @@ def process_controversial_clusters(
     return results, discussion_history, updated_consensus_proportion, updated_entropy
 
 
-def extract_consensus_metrics_from_discussion(
-    discussion: str,
-) -> tuple[Optional[float], Optional[float]]:
-    """Extract consensus proportion (CP) and entropy (H) values from discussion text.
-
-    Args:
-        discussion: Text of the model discussion
-
-    Returns:
-        tuple[Optional[float], Optional[float]]: Extracted CP and H values, or None if not found
-
-    """
-    cp, h, _ = _extract_metrics_from_text(discussion)
-    return cp, h
-
-
 def extract_cell_type_from_discussion(discussion: str) -> Optional[str]:
     """Extract the final cell type determination from a discussion.
 
@@ -1001,7 +983,6 @@ def interactive_consensus_annotation(
             raise ValueError(error_msg)
 
         # Filter marker_genes to only include specified clusters
-        original_marker_genes = marker_genes.copy()
         marker_genes = {cluster_id: marker_genes[cluster_id] for cluster_id in valid_clusters}
 
         # Log the filtering

python/mllmcelltype/providers/grok.py

@@ -45,9 +45,6 @@ def process_grok(
 
     write_log(f"Using model: {model}")
 
-    # Process all input at once
-    write_log("Processing input in 1 chunk", level="debug")
-
     # Prepare the request body
     body = {"model": model, "messages": [{"role": "user", "content": prompt}]}
 
@@ -85,8 +82,6 @@ def process_grok(
             write_log(f"Got response with {len(res)} lines")
             write_log(f"Raw response from Grok:\n{res}", level="debug")
 
-            # Success, exit retry loop
-            write_log("All chunks processed successfully", level="debug")
             # Clean up results (remove commas at the end of lines)
             return [line.rstrip(",") for line in res]
 

python/mllmcelltype/utils.py

@@ -321,23 +321,12 @@ def format_results(results: list[str], clusters: list[str]) -> dict[str, str]:
         # Extract annotations from JSON structure
         if "annotations" in data and isinstance(data["annotations"], list):
             json_result = {}
-            metadata = {}
 
             for annotation in data["annotations"]:
                 if "cluster" in annotation and "cell_type" in annotation:
                     cluster_id = annotation["cluster"]
                     json_result[cluster_id] = annotation["cell_type"]
 
-                    # Store additional metadata if available
-                    cluster_metadata = {}
-                    if "confidence" in annotation:
-                        cluster_metadata["confidence"] = annotation["confidence"]
-                    if "key_markers" in annotation:
-                        cluster_metadata["key_markers"] = annotation["key_markers"]
-
-                    if cluster_metadata:
-                        metadata[cluster_id] = cluster_metadata
-
             # If we found annotations for all clusters, return the result
             if len(json_result) == len(clusters):
                 write_log("Successfully parsed JSON response", level="info")
@@ -699,7 +688,7 @@ def get_cache_stats(cache_dir: Optional[str] = None) -> dict[str, Any]:
     oldest = float("inf")
     newest = 0
     provider_counts = {}
-    format_counts = {"legacy": 0, "v1.0": 0, "unknown": 0}
+    format_counts = {"legacy": 0, "1.0": 0, "unknown": 0}
     valid_files = 0
     invalid_files = 0