vouchervision/utils_VoucherVision_parallel.py

import openai
import os, json, glob, shutil, yaml, torch, logging
import openpyxl
from openpyxl import Workbook, load_workbook
from tqdm import tqdm
import vertexai
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from langchain_openai import AzureChatOpenAI
from google.oauth2 import service_account
from transformers import AutoTokenizer, AutoModel

from concurrent.futures import ThreadPoolExecutor, as_completed
from queue import Queue
import threading

from vouchervision.LLM_OpenAI import OpenAIHandler
from vouchervision.LLM_GooglePalm2 import GooglePalm2Handler
from vouchervision.LLM_GoogleGemini import GoogleGeminiHandler
from vouchervision.LLM_MistralAI import MistralHandler
from vouchervision.LLM_local_cpu_MistralAI import LocalCPUMistralHandler
from vouchervision.LLM_local_MistralAI import LocalMistralHandler 
from vouchervision.prompt_catalog import PromptCatalog
from vouchervision.model_maps import ModelMaps
from vouchervision.general_utils import get_cfg_from_full_path
from vouchervision.OCR_google_cloud_vision import OCREngine 

'''
* For the prefix_removal, the image names have 'MICH-V-' prior to the barcode, so that is used for matching
  but removed for output.
* There is also code active to replace the LLM-predicted "Catalog Number" with the correct number since it is known.
  The LLMs to usually assign the barcode to the correct field, but it's not needed since it is already known.
        - Look for ####################### Catalog Number pre-defined
'''


    
class VoucherVision():

    def __init__(self, cfg, logger, dir_home, path_custom_prompts, Project, Dirs, is_hf, config_vals_for_permutation=None):
        self.cfg = cfg
        self.logger = logger
        self.dir_home = dir_home
        self.path_custom_prompts = path_custom_prompts
        self.Project = Project
        self.Dirs = Dirs
        self.headers = None
        self.prompt_version = None
        self.is_hf = is_hf

        ### config_vals_for_permutation allows you to set the starting temp, top_k, top_p, seed....
        self.config_vals_for_permutation = config_vals_for_permutation

        # self.trOCR_model_version = "microsoft/trocr-large-handwritten"
        # self.trOCR_model_version = "microsoft/trocr-base-handwritten"
        # self.trOCR_model_version = "dh-unibe/trocr-medieval-escriptmask" # NOPE
        # self.trOCR_model_version = "dh-unibe/trocr-kurrent" # NOPE
        # self.trOCR_model_version = "DunnBC22/trocr-base-handwritten-OCR-handwriting_recognition_v2" # NOPE
        self.trOCR_processor = None
        self.trOCR_model = None

        self.set_API_keys()
        self.setup()


    def setup(self):
        self.logger.name = f'[Transcription]'
        self.logger.info(f'Setting up OCR and LLM')

        self.trOCR_model_version = self.cfg['leafmachine']['project']['trOCR_model_path']

        self.db_name = self.cfg['leafmachine']['project']['embeddings_database_name']
        self.path_domain_knowledge = self.cfg['leafmachine']['project']['path_to_domain_knowledge_xlsx']
        self.build_new_db = self.cfg['leafmachine']['project']['build_new_embeddings_database']

        self.continue_run_from_partial_xlsx = self.cfg['leafmachine']['project']['continue_run_from_partial_xlsx']

        self.prefix_removal = self.cfg['leafmachine']['project']['prefix_removal']
        self.suffix_removal = self.cfg['leafmachine']['project']['suffix_removal']
        self.catalog_numerical_only = self.cfg['leafmachine']['project']['catalog_numerical_only']

        self.prompt_version0 = self.cfg['leafmachine']['project']['prompt_version']
        self.use_domain_knowledge = self.cfg['leafmachine']['project']['use_domain_knowledge']

        self.catalog_name_options = ["Catalog Number", "catalog_number", "catalogNumber"]

        self.geo_headers = ["GEO_override_OCR", "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat",
                       "GEO_decimal_long","GEO_city", "GEO_county", "GEO_state",
                       "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent",]
        
        self.usage_headers = ["current_time", "inference_time_s", "tool_time_s","max_cpu", "max_ram_gb", "n_gpus", "max_gpu_load", "max_gpu_vram_gb","total_gpu_vram_gb","capability_score",]
        
        self.wfo_headers = ["WFO_override_OCR", "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement"]
        self.wfo_headers_no_lists = ["WFO_override_OCR", "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_placement"]
        
        self.utility_headers = ["filename"] + self.wfo_headers + self.geo_headers + self.usage_headers + ["run_name", "prompt", "LLM", "tokens_in", "tokens_out", "LM2_collage", "OCR_method", "OCR_double", "OCR_trOCR", "path_to_crop","path_to_original","path_to_content","path_to_helper",]
                                # "WFO_override_OCR", "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement",
                                
                                # "GEO_override_OCR", "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat",
                                # "GEO_decimal_long","GEO_city", "GEO_county", "GEO_state",
                                # "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent",
                                
                                # "tokens_in", "tokens_out", "path_to_crop","path_to_original","path_to_content","path_to_helper",]
        
        # WFO_candidate_names is separate, bc it may be type --> list

        self.do_create_OCR_helper_image = self.cfg['leafmachine']['do_create_OCR_helper_image']

        self.map_prompt_versions()
        self.map_dir_labels()
        self.map_API_options()
        # self.init_embeddings()
        self.init_transcription_xlsx()
        self.init_trOCR_model()

        '''Logging'''
        self.logger.info(f'Transcribing dataset --- {self.dir_labels}')
        self.logger.info(f'Saving transcription batch to --- {self.path_transcription}')
        self.logger.info(f'Saving individual transcription files to --- {self.Dirs.transcription_ind}')
        self.logger.info(f'Starting transcription...')
        self.logger.info(f'     LLM MODEL --> {self.version_name}')
        self.logger.info(f'     Using Azure API --> {self.is_azure}')
        self.logger.info(f'     Model name passed to API --> {self.model_name}')
        self.logger.info(f'     API access token is found in PRIVATE_DATA.yaml --> {self.has_key}')


    def init_trOCR_model(self):
        lgr = logging.getLogger('transformers')
        lgr.setLevel(logging.ERROR)
        
        self.trOCR_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") # usually just the "microsoft/trocr-base-handwritten"
        self.trOCR_model = VisionEncoderDecoderModel.from_pretrained(self.trOCR_model_version) # This matches the model
        
        # Check for GPU availability
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.trOCR_model.to(self.device)


    def map_API_options(self):
        self.chat_version = self.cfg['leafmachine']['LLM_version']

        # Get the required values from ModelMaps
        self.model_name = ModelMaps.get_version_mapping_cost(self.chat_version)
        self.is_azure = ModelMaps.get_version_mapping_is_azure(self.chat_version)
        self.has_key = ModelMaps.get_version_has_key(self.chat_version, self.has_key_openai, self.has_key_azure_openai, self.has_key_google_application_credentials, self.has_key_mistral)

        # Check if the version is supported
        if self.model_name is None:
            supported_LLMs = ", ".join(ModelMaps.get_models_gui_list())
            raise Exception(f"Unsupported LLM: {self.chat_version}. Requires one of: {supported_LLMs}")

        self.version_name = self.chat_version


    def map_prompt_versions(self):
        self.prompt_version_map = {
            "Version 1": "prompt_v1_verbose",
        }
        self.prompt_version = self.prompt_version_map.get(self.prompt_version0, self.path_custom_prompts)
        self.is_predefined_prompt = self.is_in_prompt_version_map(self.prompt_version)


    def is_in_prompt_version_map(self, value):
        return value in self.prompt_version_map.values()


    def map_dir_labels(self):
        if self.cfg['leafmachine']['use_RGB_label_images']:
            self.dir_labels = os.path.join(self.Dirs.save_per_annotation_class,'label')
        else:
            self.dir_labels = self.Dirs.save_original

        # Use glob to get all image paths in the directory
        self.img_paths = glob.glob(os.path.join(self.dir_labels, "*"))


    def load_rules_config(self):
        with open(self.path_custom_prompts, 'r') as stream:
            try:
                return yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                print(exc)
                return None
            

    def generate_xlsx_headers(self):
        # Extract headers from the 'Dictionary' keys in the JSON template rules
        # xlsx_headers = list(self.rules_config_json['rules']["Dictionary"].keys())
        xlsx_headers = list(self.rules_config_json['rules'].keys())
        xlsx_headers = xlsx_headers + self.utility_headers
        return xlsx_headers


    def init_transcription_xlsx(self):
        # Initialize output file
        self.path_transcription = os.path.join(self.Dirs.transcription,"transcribed.xlsx")
        
        # else:
        if not self.is_predefined_prompt:
            # Load the rules configuration
            self.rules_config_json = self.load_rules_config()
            # Generate the headers from the configuration
            self.headers = self.generate_xlsx_headers()
            # Set the headers used to the dynamically generated headers
            self.headers_used = 'CUSTOM'
        else:
            # If it's a predefined prompt, raise an exception as we don't have further instructions
            raise ValueError("Predefined prompt is not handled in this context.")

        self.create_or_load_excel_with_headers(os.path.join(self.Dirs.transcription,"transcribed.xlsx"), self.headers)

           
    def create_or_load_excel_with_headers(self, file_path, headers, show_head=False):
        output_dir_names = ['Archival_Components', 'Config_File', 'Cropped_Images', 'Logs', 'Original_Images', 'Transcription']
        self.completed_specimens = []

        # Check if the file exists and it's not None
        if self.continue_run_from_partial_xlsx is not None and os.path.isfile(self.continue_run_from_partial_xlsx):
            workbook = load_workbook(filename=self.continue_run_from_partial_xlsx)
            sheet = workbook.active
            show_head=True
            # Identify the 'path_to_crop' column
            try:
                path_to_crop_col = headers.index('path_to_crop') + 1
                path_to_original_col = headers.index('path_to_original') + 1
                path_to_content_col = headers.index('path_to_content') + 1
                path_to_helper_col = headers.index('path_to_helper') + 1
                # self.completed_specimens = list(sheet.iter_cols(min_col=path_to_crop_col, max_col=path_to_crop_col, values_only=True, min_row=2))
            except ValueError:
                print("'path_to_crop' not found in the header row.")

            path_to_crop = list(sheet.iter_cols(min_col=path_to_crop_col, max_col=path_to_crop_col, values_only=True, min_row=2))
            path_to_original = list(sheet.iter_cols(min_col=path_to_original_col, max_col=path_to_original_col, values_only=True, min_row=2))
            path_to_content = list(sheet.iter_cols(min_col=path_to_content_col, max_col=path_to_content_col, values_only=True, min_row=2))
            path_to_helper = list(sheet.iter_cols(min_col=path_to_helper_col, max_col=path_to_helper_col, values_only=True, min_row=2))
            others = [path_to_crop_col, path_to_original_col, path_to_content_col, path_to_helper_col]
            jsons = [path_to_content_col, path_to_helper_col]

            for cell in path_to_crop[0]:
                old_path = cell
                new_path = file_path
                for dir_name in output_dir_names:
                    if dir_name in old_path:
                        old_path_parts = old_path.split(dir_name)
                        new_path_parts = new_path.split('Transcription')
                        updated_path = new_path_parts[0] + dir_name + old_path_parts[1]
                        self.completed_specimens.append(os.path.basename(updated_path))
            print(f"{len(self.completed_specimens)} images are already completed")

            ### Copy the JSON files over
            for colu in jsons:
                cell = next(sheet.iter_rows(min_row=2, min_col=colu, max_col=colu))[0]
                old_path = cell.value
                new_path = file_path

                old_path_parts = old_path.split('Transcription')
                new_path_parts = new_path.split('Transcription')
                updated_path = new_path_parts[0] + 'Transcription' + old_path_parts[1]

                # Copy files
                old_dir = os.path.dirname(old_path)
                new_dir = os.path.dirname(updated_path)

                # Check if old_dir exists and it's a directory
                if os.path.exists(old_dir) and os.path.isdir(old_dir):
                    # Check if new_dir exists. If not, create it.
                    if not os.path.exists(new_dir):
                        os.makedirs(new_dir)

                    # Iterate through all files in old_dir and copy each to new_dir
                    for filename in os.listdir(old_dir):
                        shutil.copy2(os.path.join(old_dir, filename), new_dir) # copy2 preserves metadata

            ### Update the file names
            for colu in others:
                for row in sheet.iter_rows(min_row=2, min_col=colu, max_col=colu):
                    for cell in row:
                        old_path = cell.value
                        new_path = file_path
                        for dir_name in output_dir_names:
                            if dir_name in old_path:
                                old_path_parts = old_path.split(dir_name)
                                new_path_parts = new_path.split('Transcription')
                                updated_path = new_path_parts[0] + dir_name + old_path_parts[1]
                                cell.value = updated_path
            show_head=True

                
        else:
            # Create a new workbook and select the active worksheet
            workbook = Workbook()
            sheet = workbook.active

            # Write headers in the first row
            for i, header in enumerate(headers, start=1):
                sheet.cell(row=1, column=i, value=header)
            self.completed_specimens = []
            
        # Save the workbook
        workbook.save(file_path)

        if show_head:
            print("continue_run_from_partial_xlsx:")
            for i, row in enumerate(sheet.iter_rows(values_only=True)):
                print(row)
                if i == 3:  # print the first 5 rows (0-indexed)
                    print("\n")
                    break


    def add_data_to_excel_from_response(self, Dirs, path_transcription, response, WFO_record, GEO_record, usage_report, 
                                        MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, path_to_content, path_to_helper, nt_in, nt_out):
        

        wb = openpyxl.load_workbook(path_transcription)
        sheet = wb.active

        # find the next empty row
        next_row = sheet.max_row + 1

        if isinstance(response, str):
            try:
                response = json.loads(response)
            except json.JSONDecodeError:
                print(f"Failed to parse response: {response}")
                return

        # iterate over headers in the first row
        for i, header in enumerate(sheet[1], start=1):
            # check if header value is in response keys
            if (header.value in response) and (header.value not in self.catalog_name_options): ####################### Catalog Number pre-defined
                # check if the response value is a dictionary
                if isinstance(response[header.value], dict):
                    # if it is a dictionary, extract the 'value' field
                    cell_value = response[header.value].get('value', '')
                else:
                    # if it's not a dictionary, use it directly
                    cell_value = response[header.value]
                
                try:
                    # write the value to the cell
                    sheet.cell(row=next_row, column=i, value=cell_value)
                except:
                    sheet.cell(row=next_row, column=i, value=cell_value[0])

            elif header.value in self.catalog_name_options: 
                # if self.prefix_removal:
                #     filename_without_extension = filename_without_extension.replace(self.prefix_removal, "")
                # if self.suffix_removal:
                #     filename_without_extension = filename_without_extension.replace(self.suffix_removal, "")
                # if self.catalog_numerical_only:
                #     filename_without_extension = self.remove_non_numbers(filename_without_extension)
                sheet.cell(row=next_row, column=i, value=filename_without_extension)
            elif header.value == "path_to_crop":
                sheet.cell(row=next_row, column=i, value=path_to_crop)
            elif header.value == "path_to_original":
                if self.cfg['leafmachine']['use_RGB_label_images']:
                    fname = os.path.basename(path_to_crop)
                    base = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(path_to_crop))))
                    path_to_original = os.path.join(base, 'Original_Images', fname)
                    sheet.cell(row=next_row, column=i, value=path_to_original)
                else:
                    fname = os.path.basename(path_to_crop)
                    base = os.path.dirname(os.path.dirname(path_to_crop))
                    path_to_original = os.path.join(base, 'Original_Images', fname)
                    sheet.cell(row=next_row, column=i, value=path_to_original)
            elif header.value == "path_to_content":
                sheet.cell(row=next_row, column=i, value=path_to_content)
            elif header.value == "path_to_helper":
                sheet.cell(row=next_row, column=i, value=path_to_helper)
            elif header.value == "tokens_in":
                sheet.cell(row=next_row, column=i, value=nt_in)
            elif header.value == "tokens_out":
                sheet.cell(row=next_row, column=i, value=nt_out)
            elif header.value == "filename":
                sheet.cell(row=next_row, column=i, value=filename_without_extension)
            elif header.value == "prompt":
                sheet.cell(row=next_row, column=i, value=os.path.basename(self.path_custom_prompts))
            elif header.value == "run_name":
                sheet.cell(row=next_row, column=i, value=Dirs.run_name)
            elif header.value == "LM2_collage":
                sheet.cell(row=next_row, column=i, value=self.cfg['leafmachine']['use_RGB_label_images'])
            elif header.value == "OCR_method":
                value_to_insert = self.cfg['leafmachine']['project']['OCR_option']
                if isinstance(value_to_insert, list):
                    value_to_insert = '|'.join(map(str, value_to_insert))
                sheet.cell(row=next_row, column=i, value=value_to_insert)
            elif header.value == "OCR_double":
                sheet.cell(row=next_row, column=i, value=self.cfg['leafmachine']['project']['double_OCR'])
            elif header.value == "OCR_trOCR":
                sheet.cell(row=next_row, column=i, value=self.cfg['leafmachine']['project']['do_use_trOCR'])
            # "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement"
            elif header.value in self.wfo_headers_no_lists:
                sheet.cell(row=next_row, column=i, value=WFO_record.get(header.value, ''))
            # elif header.value == "WFO_exact_match":
            #     sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_exact_match",''))
            # elif header.value == "WFO_exact_match_name":
            #     sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_exact_match_name",''))
            # elif header.value == "WFO_best_match":
            #     sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_best_match",''))
            # elif header.value == "WFO_placement":
            #     sheet.cell(row=next_row, column=i, value= WFO_record.get("WFO_placement",''))
            elif header.value == "WFO_candidate_names":
                candidate_names = WFO_record.get("WFO_candidate_names", '')
                # Check if candidate_names is a list and convert to a string if it is
                if isinstance(candidate_names, list):
                    candidate_names_str = '|'.join(candidate_names)
                else:
                    candidate_names_str = candidate_names
                sheet.cell(row=next_row, column=i, value=candidate_names_str)
            
            # "GEO_method", "GEO_formatted_full_string", "GEO_decimal_lat", "GEO_decimal_long",
            # "GEO_city", "GEO_county", "GEO_state", "GEO_state_code", "GEO_country", "GEO_country_code", "GEO_continent"
            elif header.value in self.geo_headers:
                sheet.cell(row=next_row, column=i, value=GEO_record.get(header.value, ''))

            elif header.value in self.usage_headers:
                sheet.cell(row=next_row, column=i, value=usage_report.get(header.value, ''))

            elif header.value == "LLM":
                sheet.cell(row=next_row, column=i, value=MODEL_NAME_FORMATTED)

        # save the workbook
        wb.save(path_transcription)
    

    def has_API_key(self, val):
        return isinstance(val, str) and bool(val.strip())
        # if val != '':
        #     return True
        # else:
        #     return False
        

    def get_google_credentials(self): # Also used for google drive
        if self.is_hf:
            creds_json_str = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')
            credentials = service_account.Credentials.from_service_account_info(json.loads(creds_json_str))
            return credentials
        else:
            with open(self.cfg_private['google']['GOOGLE_APPLICATION_CREDENTIALS'], 'r') as file:
                data = json.load(file)
                creds_json_str = json.dumps(data)
                credentials = service_account.Credentials.from_service_account_info(json.loads(creds_json_str))
                os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = creds_json_str
                return credentials
        

    def set_API_keys(self):
        if self.is_hf:
            self.dir_home = os.path.dirname(os.path.dirname(__file__))
            self.path_cfg_private = None
            self.cfg_private = None

            k_openai = os.getenv('OPENAI_API_KEY')
            k_openai_azure = os.getenv('AZURE_API_VERSION')

            k_google_project_id = os.getenv('GOOGLE_PROJECT_ID')
            k_google_location = os.getenv('GOOGLE_LOCATION')
            k_google_application_credentials = os.getenv('GOOGLE_APPLICATION_CREDENTIALS')

            k_mistral = os.getenv('MISTRAL_API_KEY')
            k_here = os.getenv('HERE_API_KEY')
            k_opencage = os.getenv('open_cage_geocode')
        else:
            self.dir_home = os.path.dirname(os.path.dirname(__file__))
            self.path_cfg_private = os.path.join(self.dir_home, 'PRIVATE_DATA.yaml')
            self.cfg_private = get_cfg_from_full_path(self.path_cfg_private)

            k_openai = self.cfg_private['openai']['OPENAI_API_KEY']
            k_openai_azure = self.cfg_private['openai_azure']['OPENAI_API_KEY_AZURE']

            k_google_project_id = self.cfg_private['google']['GOOGLE_PROJECT_ID']
            k_google_location = self.cfg_private['google']['GOOGLE_LOCATION']
            k_google_application_credentials = self.cfg_private['google']['GOOGLE_APPLICATION_CREDENTIALS']
            
            k_mistral = self.cfg_private['mistral']['MISTRAL_API_KEY']
            k_here = self.cfg_private['here']['API_KEY']
            k_opencage = self.cfg_private['open_cage_geocode']['API_KEY']
            


        self.has_key_openai = self.has_API_key(k_openai)
        self.has_key_azure_openai = self.has_API_key(k_openai_azure)
        self.llm = None
        
        self.has_key_google_project_id = self.has_API_key(k_google_project_id)
        self.has_key_google_location = self.has_API_key(k_google_location)
        self.has_key_google_application_credentials = self.has_API_key(k_google_application_credentials)

        self.has_key_mistral = self.has_API_key(k_mistral)
        self.has_key_here = self.has_API_key(k_here)
        self.has_key_open_cage_geocode = self.has_API_key(k_opencage)

        

        ### Google - OCR, Palm2, Gemini
        if self.has_key_google_application_credentials and self.has_key_google_project_id and self.has_key_google_location:
            if self.is_hf:
                vertexai.init(project=os.getenv('GOOGLE_PROJECT_ID'), location=os.getenv('GOOGLE_LOCATION'), credentials=self.get_google_credentials())
            else:
                vertexai.init(project=k_google_project_id, location=k_google_location, credentials=self.get_google_credentials())
                os.environ['GOOGLE_API_KEY'] = self.cfg_private['google']['GOOGLE_PALM_API']


        ### OpenAI
        if self.has_key_openai:
            if self.is_hf:
                openai.api_key = os.getenv('OPENAI_API_KEY')
            else:
                openai.api_key = self.cfg_private['openai']['OPENAI_API_KEY']
                os.environ["OPENAI_API_KEY"] = self.cfg_private['openai']['OPENAI_API_KEY']


        ### OpenAI - Azure
        if self.has_key_azure_openai:
            if self.is_hf:
                # Initialize the Azure OpenAI client
                self.llm = AzureChatOpenAI(
                    deployment_name = 'gpt-35-turbo',#'gpt-35-turbo',
                    openai_api_version = os.getenv('AZURE_API_VERSION'),
                    openai_api_key = os.getenv('AZURE_API_KEY'),
                    azure_endpoint = os.getenv('AZURE_API_BASE'),
                    openai_organization = os.getenv('AZURE_ORGANIZATION'),
                )
                
            else:
                # Initialize the Azure OpenAI client
                self.llm = AzureChatOpenAI(
                    deployment_name = 'gpt-35-turbo',#'gpt-35-turbo',
                    openai_api_version = self.cfg_private['openai_azure']['OPENAI_API_VERSION'],
                    openai_api_key = self.cfg_private['openai_azure']['OPENAI_API_KEY_AZURE'],
                    azure_endpoint = self.cfg_private['openai_azure']['OPENAI_API_BASE'],
                    openai_organization = self.cfg_private['openai_azure']['OPENAI_ORGANIZATION'],
                )
                

        ### Mistral
        if self.has_key_mistral:
            if self.is_hf:
                pass # Already set
            else:
                os.environ['MISTRAL_API_KEY'] = self.cfg_private['mistral']['MISTRAL_API_KEY']


        ### HERE
        if self.has_key_here:
            if self.is_hf:
                pass # Already set
            else:
                os.environ['HERE_APP_ID'] = self.cfg_private['here']['APP_ID']
                os.environ['HERE_API_KEY'] = self.cfg_private['here']['API_KEY']


        ### HERE
        if self.has_key_open_cage_geocode:
            if self.is_hf:
                pass # Already set
            else:
                os.environ['OPENCAGE_API_KEY'] = self.cfg_private['open_cage_geocode']['API_KEY']
                

        
    def clean_catalog_number(self, data, filename_without_extension):
        #Cleans up the catalog number in data if it's a dict
        
        def modify_catalog_key(catalog_key, filename_without_extension, data):
            # Helper function to apply modifications on catalog number
            if catalog_key not in data:
                new_data = {catalog_key: None}
                data = {**new_data, **data}

            if self.prefix_removal:
                filename_without_extension = filename_without_extension.replace(self.prefix_removal, "")
            if self.suffix_removal:
                filename_without_extension = filename_without_extension.replace(self.suffix_removal, "")
            if self.catalog_numerical_only:
                filename_without_extension = self.remove_non_numbers(data[catalog_key])
            data[catalog_key] = filename_without_extension
            return data
        
        if isinstance(data, dict):
            if self.headers_used == 'HEADERS_v1_n22':
                return modify_catalog_key("Catalog Number", filename_without_extension, data)
            elif self.headers_used in ['HEADERS_v2_n26', 'CUSTOM']:
                return modify_catalog_key("filename", filename_without_extension, data)
            else:
                raise ValueError("Invalid headers used.")
        else:
            raise TypeError("Data is not of type dict.")
        

    def write_json_to_file(self, filepath, data):
        '''Writes dictionary data to a JSON file.'''
        with open(filepath, 'w') as txt_file:
            if isinstance(data, dict):
                data = json.dumps(data, indent=4, sort_keys=False)
            txt_file.write(data)


    # def create_null_json(self):
    #     return {}
    

    def remove_non_numbers(self, s):
        return ''.join([char for char in s if char.isdigit()])
    

    def create_null_row(self, filename_without_extension, path_to_crop, path_to_content, path_to_helper):
        json_dict = {header: '' for header in self.headers} 
        for header, value in json_dict.items():
            if header == "path_to_crop":
                json_dict[header] = path_to_crop
            elif header == "path_to_original":
                fname = os.path.basename(path_to_crop)
                base = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(path_to_crop))))
                path_to_original = os.path.join(base, 'Original_Images', fname)
                json_dict[header] = path_to_original
            elif header == "path_to_content":
                json_dict[header] = path_to_content
            elif header == "path_to_helper":
                json_dict[header] = path_to_helper
            elif header == "filename":
                json_dict[header] = filename_without_extension

            # "WFO_exact_match","WFO_exact_match_name","WFO_best_match","WFO_candidate_names","WFO_placement"
            elif header == "WFO_exact_match":
                json_dict[header] =''
            elif header == "WFO_exact_match_name":
                json_dict[header] = ''
            elif header == "WFO_best_match":
                json_dict[header] = ''
            elif header == "WFO_candidate_names":
                json_dict[header] = ''
            elif header == "WFO_placement":
                json_dict[header] = ''
        return json_dict
    

    ##################################################################################################################################
    ##################################################     OCR      ##################################################################
    ##################################################################################################################################
    def perform_OCR_and_save_results(self, image_index, json_report, jpg_file_path_OCR_helper, txt_file_path_OCR, txt_file_path_OCR_bounds):
        self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Starting OCR')
        # self.OCR - None

        ### Process_image() runs the OCR for text, handwriting, trOCR AND creates the overlay image
        ocr_google = OCREngine(self.logger, json_report, self.dir_home, self.is_hf, self.path_to_crop, self.cfg, self.trOCR_model_version, self.trOCR_model, self.trOCR_processor, self.device)  
        ocr_google.process_image(self.do_create_OCR_helper_image, self.logger)
        self.OCR = ocr_google.OCR
        self.logger.info(f"Complete OCR text for LLM prompt:\n\n{self.OCR}\n\n")

        self.write_json_to_file(txt_file_path_OCR, ocr_google.OCR_JSON_to_file)
        
        self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Finished OCR')

        if len(self.OCR) > 0:
            ocr_google.overlay_image.save(jpg_file_path_OCR_helper)

            OCR_bounds = {}
            if ocr_google.hand_text_to_box_mapping is not None:
                OCR_bounds['OCR_bounds_handwritten'] = ocr_google.hand_text_to_box_mapping

            if ocr_google.normal_text_to_box_mapping is not None:
                OCR_bounds['OCR_bounds_printed'] = ocr_google.normal_text_to_box_mapping

            if ocr_google.trOCR_text_to_box_mapping is not None:
                OCR_bounds['OCR_bounds_trOCR'] = ocr_google.trOCR_text_to_box_mapping

            self.write_json_to_file(txt_file_path_OCR_bounds, OCR_bounds)
            self.logger.info(f'Working on {image_index + 1}/{len(self.img_paths)} --- Saved OCR Overlay Image')
        else:
            pass ########################################################################################################################### fix logic for no OCR

    ##################################################################################################################################
    #######################################################  LLM Switchboard  ########################################################
    ##################################################################################################################################
    def send_to_LLM(self, is_azure, progress_report, json_report, model_name):
        self.n_failed_LLM_calls = 0
        self.n_failed_OCR = 0

        final_JSON_response = None
        final_WFO_record = None
        final_GEO_record = None

        self.initialize_token_counters()
        self.update_progress_report_initial(progress_report)

        MODEL_NAME_FORMATTED = ModelMaps.get_API_name(model_name)
        name_parts = model_name.split("_")
        
        self.setup_JSON_dict_structure()

        Copy_Prompt = PromptCatalog()
        Copy_Prompt.copy_prompt_template_to_new_dir(self.Dirs.transcription_prompt, self.path_custom_prompts)
        
        if json_report:
            json_report.set_text(text_main=f'Loading {MODEL_NAME_FORMATTED}')
            json_report.set_JSON({}, {}, {})
        # llm_model = self.initialize_llm_model(self.cfg, self.logger, MODEL_NAME_FORMATTED, self.JSON_dict_structure, name_parts, is_azure, self.llm, self.config_vals_for_permutation)

        results_queue = Queue()
        
        if json_report:
            json_report.set_text(text_main='Sending batch to OCR and LLM')

        num_files = len(self.img_paths)
        # num_threads = min(num_files, 128)
        num_threads = 128
        counter = AtomicCounter()

        # Setup for parallel execution
        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(self.send_to_LLM_worker, 
                                    path_to_crop, 
                                    results_queue, 
                                    model_name,
                                    MODEL_NAME_FORMATTED,
                                    name_parts,
                                    is_azure, 
                                    i
                                ) for i, path_to_crop in enumerate(self.img_paths)]
            for future in tqdm(as_completed(futures), total=len(futures), desc="Processing", unit="task"):
                try:
                    # Here, you could also directly process results if they were not being put in a queue
                    future.result()  # Forces a wait on the future and re-raises any exceptions
                    new_value = counter.inc()
                    try:
                        if json_report:
                            current_value = counter.value
                            json_report.set_text(text_main=f'Completed {current_value} of {num_files}')
                    except:
                        pass
                except Exception as e:
                    # Log the error, possibly mark the task for retry, or handle it as necessary
                    print(f"A task failed with exception: {e}")

        # Process results from the queue
        while not results_queue.empty():
            response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report, path_to_crop, paths = results_queue.get()

            self.n_failed_LLM_calls += 1 if response_candidate is None else 0
                
            ### Estimate n tokens returned
            self.logger.info(f'Prompt tokens IN --- {nt_in}')
            self.logger.info(f'Prompt tokens OUT --- {nt_out}')
                
            self.update_token_counters(nt_in, nt_out)

            final_JSON_response, final_WFO_record, final_GEO_record = self.update_final_response(response_candidate, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, paths, path_to_crop, nt_in, nt_out)

            self.logger.info(f'Finished LLM call')
            
            if json_report:
                json_report.set_JSON(final_JSON_response, final_WFO_record, final_GEO_record)

        if json_report:
            json_report.set_text(text_main='Finished!')

        self.update_progress_report_final(progress_report)
        final_JSON_response = self.parse_final_json_response(final_JSON_response)
        return final_JSON_response, final_WFO_record, final_GEO_record, self.total_tokens_in, self.total_tokens_out
    
    def send_to_LLM_worker(self, path_to_crop, queue, model_name, MODEL_NAME_FORMATTED, name_parts, is_azure, i):
        llm_model = self.initialize_llm_model(self.cfg, self.logger, MODEL_NAME_FORMATTED, self.JSON_dict_structure, name_parts, is_azure, self.llm, self.config_vals_for_permutation)

        # self.update_progress_report_batch(progress_report, i)

        if self.should_skip_specimen(path_to_crop):
            self.log_skipping_specimen(path_to_crop)
            return

        paths = self.generate_paths(path_to_crop, i)
        self.path_to_crop = path_to_crop

        filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper, json_file_path_wiki, txt_file_path_ind_prompt = paths
        # if json_report:
            # json_report.set_text(text_main='Starting OCR')
        self.perform_OCR_and_save_results(i, None, jpg_file_path_OCR_helper, txt_file_path_OCR, txt_file_path_OCR_bounds)
        # if json_report:
            # json_report.set_text(text_main='Finished OCR')

        if not self.OCR:
            self.n_failed_OCR += 1
            response_candidate = None
            nt_in = 0
            nt_out = 0
        else:
            ### Format prompt
            prompt = self.setup_prompt()
            # prompt = remove_colons_and_double_apostrophes(prompt) # This is moved to utils_VV since it broke the json structure.

            ### Send prompt to chosen LLM
            self.logger.info(f'Waiting for {model_name} API call --- Using {MODEL_NAME_FORMATTED}')

            if 'PALM2' in name_parts:
                response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_api_GooglePalm2(prompt, None, paths)
            
            elif 'GEMINI' in name_parts:
                response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_api_GoogleGemini(prompt, None, paths)
            
            elif 'MISTRAL' in name_parts and ('LOCAL' not in name_parts):
                response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_api_MistralAI(prompt, None, paths)
            
            elif 'LOCAL' in name_parts: 
                if 'MISTRAL' in name_parts or 'MIXTRAL' in name_parts:
                    if 'CPU' in name_parts:     
                        response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_local_cpu_MistralAI(prompt, None, paths) 
                    else:
                        response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_local_MistralAI(prompt, None, paths) 
            else:
                response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report = llm_model.call_llm_api_OpenAI(prompt, None, paths)

        # Instead of directly updating shared resources, put the structured result in the queue
        queue.put((response_candidate, nt_in, nt_out, WFO_record, GEO_record, usage_report, path_to_crop, paths))

    ##################################################################################################################################
    ################################################## LLM Helper Funcs ##############################################################
    ##################################################################################################################################
    def initialize_llm_model(self, cfg, logger, model_name, JSON_dict_structure, name_parts, is_azure=None, llm_object=None, config_vals_for_permutation=None):
        if 'LOCAL'in name_parts:
            if ('MIXTRAL' in name_parts) or ('MISTRAL' in name_parts):
                if 'CPU' in name_parts:
                    return LocalCPUMistralHandler(cfg, logger, model_name, JSON_dict_structure, config_vals_for_permutation)
                else:
                    return LocalMistralHandler(cfg, logger, model_name, JSON_dict_structure, config_vals_for_permutation)
        else:
            if 'PALM2' in name_parts:
                return GooglePalm2Handler(cfg, logger, model_name, JSON_dict_structure, config_vals_for_permutation)
            elif 'GEMINI' in name_parts:
                return GoogleGeminiHandler(cfg, logger, model_name, JSON_dict_structure, config_vals_for_permutation)
            elif 'MISTRAL' in name_parts and ('LOCAL' not in name_parts):
                return MistralHandler(cfg, logger, model_name, JSON_dict_structure, config_vals_for_permutation)
            else:
                return OpenAIHandler(cfg, logger, model_name, JSON_dict_structure, is_azure, llm_object, config_vals_for_permutation)

    def setup_prompt(self):
        Catalog = PromptCatalog()
        prompt, _ = Catalog.prompt_SLTP(self.path_custom_prompts, OCR=self.OCR)
        return prompt
    
    def setup_JSON_dict_structure(self):
        Catalog = PromptCatalog()
        _, self.JSON_dict_structure = Catalog.prompt_SLTP(self.path_custom_prompts, OCR='Text')
    

    def initialize_token_counters(self):
        self.total_tokens_in = 0
        self.total_tokens_out = 0


    def update_progress_report_initial(self, progress_report):
        if progress_report is not None:
            progress_report.set_n_batches(len(self.img_paths))


    def update_progress_report_batch(self, progress_report, batch_index):
        if progress_report is not None:
            progress_report.update_batch(f"Working on image {batch_index + 1} of {len(self.img_paths)}")


    def should_skip_specimen(self, path_to_crop):
        return os.path.basename(path_to_crop) in self.completed_specimens


    def log_skipping_specimen(self, path_to_crop):
        self.logger.info(f'[Skipping] specimen {os.path.basename(path_to_crop)} already processed')

    
    def update_token_counters(self, nt_in, nt_out):
        self.total_tokens_in += nt_in
        self.total_tokens_out += nt_out


    def update_final_response(self, response_candidate, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, paths, path_to_crop, nt_in, nt_out):
        filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper, json_file_path_wiki, txt_file_path_ind_prompt = paths
        # Saving the JSON and XLSX files with the response and updating the final JSON response
        if response_candidate is not None:
            final_JSON_response_updated = self.save_json_and_xlsx(self.Dirs, response_candidate, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out)
            return final_JSON_response_updated, WFO_record, GEO_record
        else:
            final_JSON_response_updated = self.save_json_and_xlsx(self.Dirs, response_candidate, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out)
            return final_JSON_response_updated, WFO_record, GEO_record


    def update_progress_report_final(self, progress_report):
        if progress_report is not None:
            progress_report.reset_batch("Batch Complete")


    def parse_final_json_response(self, final_JSON_response):
        try:
            return json.loads(final_JSON_response.strip('```').replace('json\n', '', 1).replace('json', '', 1))
        except:
            return final_JSON_response
    
    

    def generate_paths(self, path_to_crop, i):
        filename_without_extension = os.path.splitext(os.path.basename(path_to_crop))[0]
        txt_file_path = os.path.join(self.Dirs.transcription_ind, filename_without_extension + '.json')
        txt_file_path_OCR = os.path.join(self.Dirs.transcription_ind_OCR, filename_without_extension + '.json')
        txt_file_path_OCR_bounds = os.path.join(self.Dirs.transcription_ind_OCR_bounds, filename_without_extension + '.json')
        jpg_file_path_OCR_helper = os.path.join(self.Dirs.transcription_ind_OCR_helper, filename_without_extension + '.jpg')
        json_file_path_wiki = os.path.join(self.Dirs.transcription_ind_wiki, filename_without_extension + '.json')
        txt_file_path_ind_prompt = os.path.join(self.Dirs.transcription_ind_prompt, filename_without_extension + '.txt')

        self.logger.info(f'Working on {i+1}/{len(self.img_paths)} --- {filename_without_extension}')

        return filename_without_extension, txt_file_path, txt_file_path_OCR, txt_file_path_OCR_bounds, jpg_file_path_OCR_helper, json_file_path_wiki, txt_file_path_ind_prompt


    def save_json_and_xlsx(self, Dirs, response, WFO_record, GEO_record, usage_report, 
                           MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out):
        if response is None:
            response = self.JSON_dict_structure
            # Insert 'filename' as the first key
            response = {'filename': filename_without_extension, **{k: v for k, v in response.items() if k != 'filename'}}
            self.write_json_to_file(txt_file_path, response)

            # Then add the null info to the spreadsheet
            response_null = self.create_null_row(filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper)
            self.add_data_to_excel_from_response(Dirs, self.path_transcription, response_null, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in=0, nt_out=0)
        
        ### Set completed JSON
        else:
            response = self.clean_catalog_number(response, filename_without_extension)
            self.write_json_to_file(txt_file_path, response)
            # add to the xlsx file
            self.add_data_to_excel_from_response(Dirs, self.path_transcription, response, WFO_record, GEO_record, usage_report, MODEL_NAME_FORMATTED, filename_without_extension, path_to_crop, txt_file_path, jpg_file_path_OCR_helper, nt_in, nt_out)
        return response
    

    def process_specimen_batch(self, progress_report, json_report, is_real_run=False):
        if not self.has_key:
            self.logger.error(f'No API key found for {self.version_name}')
            raise Exception(f"No API key found for {self.version_name}")

        try:
            if is_real_run:
                progress_report.update_overall(f"Transcribing Labels")

            final_json_response, final_WFO_record, final_GEO_record, total_tokens_in, total_tokens_out = self.send_to_LLM(self.is_azure, progress_report, json_report, self.model_name)
            
            return final_json_response, final_WFO_record, final_GEO_record, total_tokens_in, total_tokens_out

        except Exception as e:
            self.logger.error(f"LLM call failed in process_specimen_batch: {e}")
            if progress_report is not None:
                progress_report.reset_batch(f"Batch Failed")
            self.close_logger_handlers()
            raise


    def close_logger_handlers(self):
        for handler in self.logger.handlers[:]:
            handler.close()
            self.logger.removeHandler(handler)


    # def process_specimen_batch_OCR_test(self, path_to_crop):
    #     for img_filename in os.listdir(path_to_crop):
    #         img_path = os.path.join(path_to_crop, img_filename)
    #     self.OCR, self.bounds, self.text_to_box_mapping = detect_text(img_path)

# https://gist.github.com/benhoyt/8c8a8d62debe8e5aa5340373f9c509c7
class AtomicCounter(object):
        """An atomic, thread-safe counter"""
    
        def __init__(self, initial=0):
            """Initialize a new atomic counter to given initial value"""
            self._value = initial
            self._lock = threading.Lock()
    
        def inc(self, num=1):
            """Atomically increment the counter by num and return the new value"""
            with self._lock:
                self._value += num
                return self._value
    
        def dec(self, num=1):
            """Atomically decrement the counter by num and return the new value"""
            with self._lock:
                self._value -= num
                return self._value
    
        @property
        def value(self):
            return self._value
        

def space_saver(cfg, Dirs, logger):
    dir_out = cfg['leafmachine']['project']['dir_output']
    run_name = Dirs.run_name

    path_project = os.path.join(dir_out, run_name)

    if cfg['leafmachine']['project']['delete_temps_keep_VVE']:
        logger.name = '[DELETE TEMP FILES]'
        logger.info("Deleting temporary files. Keeping files required for VoucherVisionEditor.")
        delete_dirs = ['Archival_Components', 'Config_File']
        for d in delete_dirs:
            path_delete = os.path.join(path_project, d)
            if os.path.exists(path_delete):
                shutil.rmtree(path_delete)

    elif cfg['leafmachine']['project']['delete_all_temps']:
        logger.name = '[DELETE TEMP FILES]'
        logger.info("Deleting ALL temporary files!")
        delete_dirs = ['Archival_Components', 'Config_File', 'Original_Images', 'Cropped_Images']
        for d in delete_dirs:
            path_delete = os.path.join(path_project, d)
            if os.path.exists(path_delete):
                shutil.rmtree(path_delete)

        # Delete the transctiption folder, but keep the xlsx
        transcription_path = os.path.join(path_project, 'Transcription')
        if os.path.exists(transcription_path):
            for item in os.listdir(transcription_path):
                item_path = os.path.join(transcription_path, item)
                if os.path.isdir(item_path):  # if the item is a directory
                    if os.path.exists(item_path):
                        shutil.rmtree(item_path)  # delete the directory