/ postprocessors / json_simplifier.py
json_simplifier.py
  1  """
  2  JSON Simplifier Post-processor
  3  --------------------------------
  4  This module contains a post-processor that transforms complex JSON results
  5  into clear English text using an LLM.
  6  """
  7  
  8  import json
  9  import logging
 10  from typing import Dict, Any, Optional
 11  from model_manager import ModelManager
 12  
 13  # Logging configuration
 14  logger = logging.getLogger("json_simplifier")
 15  
 16  class JSONSimplifier:
 17      """
 18      Post-processor that converts complex JSON results into clear English text.
 19      Uses an LLM model to generate a textual explanation of JSON content.
 20      """
 21      
 22      def __init__(self, config: Dict[str, Any]):
 23          """
 24          Initializes the post-processor with the specified configuration.
 25          
 26          Args:
 27              config: Post-processor configuration
 28          """
 29          self.enabled = config.get("enabled", False)
 30          self.model_name = config.get("model", "huihui-ai/DeepSeek-R1-Distill-Qwen-14B-abliterated-v2")
 31          self.system_prompt = config.get("system_prompt", "Translate this json {text} in plain english")
 32          self.max_tokens = config.get("max_tokens", 1000)
 33          self.temperature = config.get("temperature", 0.3)
 34          self.apply_to = config.get("apply_to", ["inference", "video", "transcription"])
 35          
 36          # Model initialization (loaded on demand)
 37          self.model = None
 38          
 39          logger.info(f"JSONSimplifier initialized: enabled={self.enabled}, model={self.model_name}")
 40      
 41      def should_process(self, task_type: str) -> bool:
 42          """
 43          Determines if this post-processor should be applied to the specified task type.
 44          
 45          Args:
 46              task_type: Task type (inference, video, transcription, etc.)
 47              
 48          Returns:
 49              True if the post-processor should be applied, False otherwise
 50          """
 51          return self.enabled and task_type in self.apply_to
 52      
 53      def process(self, result: Dict[str, Any], task_type: str) -> Dict[str, Any]:
 54          """
 55          Processes the JSON result and adds a plain text explanation.
 56          
 57          Args:
 58              result: JSON result to process
 59              task_type: Type of task that generated this result
 60              
 61          Returns:
 62              JSON result with the explanation added
 63          """
 64          if not self.should_process(task_type):
 65              return result
 66          
 67          try:
 68              # Check if the result contains data to process
 69              if "result" not in result or not result["result"]:
 70                  logger.warning("No result to simplify found")
 71                  return result
 72              
 73              # Load the model if necessary
 74              if self.model is None:
 75                  logger.info(f"Loading model {self.model_name} for JSON simplification")
 76                  try:
 77                      model_manager = ModelManager.get_instance()
 78                      self.model = model_manager.get_model("llm", self.model_name)
 79                  except Exception as e:
 80                      logger.error(f"Error loading model: {str(e)}")
 81                      return result
 82              
 83              # Prepare the prompt with the JSON
 84              json_str = json.dumps(result["result"], ensure_ascii=False)
 85              prompt = self.system_prompt.format(text=json_str)
 86              
 87              # Generate the explanation
 88              logger.info("Generating plain text explanation of JSON")
 89              explanation = self.model.generate(
 90                  prompt,
 91                  max_tokens=self.max_tokens,
 92                  temperature=self.temperature
 93              )
 94              
 95              # Add the explanation to the result
 96              result["plain_explanation"] = explanation.strip()
 97              
 98              logger.info("Plain text explanation generated successfully")
 99              return result
100              
101          except Exception as e:
102              logger.error(f"Error during JSON simplification: {str(e)}")
103              return result