# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_claude_corrector.ipynb.

# %% auto 0
__all__ = ['ClaudeDiarizationCorrector']

# %% ../nbs/00_claude_corrector.ipynb 6
from claudette import Chat, models
import os

class ClaudeDiarizationCorrector:
    def __init__(self, input_path: str, output_path: str, chunk_size: int = 20000):
        """
        Initialize the ClaudeDiarizationCorrector.
        
        Args:
            input_path (str): Path to the input diarization transcript file.
            output_path (str): Path where the corrected transcript will be saved.
            chunk_size (int, optional): Maximum size of each chunk for processing. Defaults to 20000.
        """
        self.input_path = input_path
        self.output_path = output_path
        self.chunk_size = chunk_size
        self.model = models[1]  # sonnet 3.5
        self.chat = Chat(self.model, sp="""You are a helpful and concise assistant that will correct a diarization script.""")

    def read_and_chunk(self):
        """
        Read the input file and split it into chunks.
        
        Returns:
            list: A list of text chunks.
        """
        with open(self.input_path, 'r', encoding='utf-8') as file:
            text = file.read()

        lines = text.split('\n')
        chunks = []
        current_chunk = ""

        for line in lines:
            if len(current_chunk + '\n' + line) > self.chunk_size:
                chunks.append(current_chunk)
                current_chunk = line
            else:
                if current_chunk:
                    current_chunk += '\n' + line
                else:
                    current_chunk = line

        if current_chunk:
            chunks.append(current_chunk)
        return chunks

    def process_chunk(self, chunk: str, index: int, total_chunks: int):
        """
        Process a single chunk of the diarization transcript.
        
        Args:
            chunk (str): The text chunk to process.
            index (int): The index of the current chunk.
            total_chunks (int): The total number of chunks.
        
        Returns:
            str: The corrected text chunk.
        """
        if index % 4 == 0:
            self.chat = Chat(self.model, sp="""You are a helpful and concise assistant that will correct a diarization script.""")
            prompt = f'''In the speaker diarization transcript below, some words are potentially mis-attributed due to bleeding or other common diarization errors. The words are always in the right order, but they may be attributed to the wrong speaker.
                      For example one speaker is talking then immediately after a second speaker is talking, and the first words of the second speaker are attributed to the last words of the second speaker. 
                      I need you to read the meaning and context of the sentences and make your best guess about how the sentences should be re-distributed 
                      They may also be displaced in the following manner-- Two or more speakers speak, and the words are all attributed to one speaker, 
                      when actually 2 or three speakers spoke like this: speaker_1 "talking" speaker_2 "talking" speaker_1 "talking", but all the transcription captures is speaker_1 "talking talking talking" 
                      Please correct the attribution of those words with the correct speaker. Directly return the corrected transcript without explaining what changes were made or why you made those changes. Use your judgement to determine the correct speaker attribution.
                      Also, if you think speaker attribution has mis-attributed in any other novel way, please correct that as well.
                      Also, if any speakers are attributes as "speaker_unknown" then make your best guess as to who the speaker is.
                      This is either the first of {total_chunks} parts, or the context has been reset in order to conserve computing costs. Correct this chunk in isolation. Text to correct is indicated by this tag <text-to-correct></text-to-correct>
                      <text-to-correct>{chunk}</text-to-correct>
                    '''
        else:
            prompt = f'''In the speaker diarization transcript below, some words are potentially mis-attributed due to bleeding or other common diarization errors. The words are always in the right order, but they may be attributed to the wrong speaker.
                      For example one speaker is talking then immediately after a second speaker is talking, and the first words of the second speaker are attributed to the last words of the second speaker. 
                      I need you to read the meaning and context of the sentences and make your best guess about how the sentences should be re-distributed 
                      They may also be displaced in the following manner-- Two or more speakers speak, and the words are all attributed to one speaker, 
                      when actually 2 or three speakers spoke like this: speaker_1 "talking" speaker_2 "talking" speaker_1 "talking", but all the transcription captures is speaker_1 "talking talking talking" 
                      Please correct the attribution of those words with the correct speaker. Directly return the corrected transcript without explaining what changes were made or why you made those changes. Use your judgement to determine the correct speaker attribution.
                      Also, if you think speaker attribution has mis-attributed in any other novel way, please correct that as well.
                      Continue diarization from previous {index} parts, maintaining consistent speaker attribution and context. The diarization transcript to correct is indicated by this tag <text-to-correct></text-to-correct>
                      Also, if any speakers are attributes as "speaker_unknown" then make your best guess as to who the speaker is.
                      This is part {index + 1} of {total_chunks}.
                      <text-to-correct>{chunk}</text-to-correct>
                    '''
        r = self.chat(prompt)
        print(f"Processed chunk {index + 1} of {total_chunks}")
        corrected_text = r.content[0].text
        return corrected_text

    def process_conversation(self):
        """
        Process the entire conversation, chunk by chunk.
        
        Returns:
            str: The fully corrected transcript.
        """
        chunks = self.read_and_chunk()
        results = []

        for index, chunk in enumerate(chunks):
            corrected_text = self.process_chunk(chunk, index, len(chunks))
            results.append(corrected_text)

        processed_text = '\n'.join(results)
        self.write_results_to_file(processed_text)
        return processed_text

    def write_results_to_file(self, processed_text: str):
        """
        Write the processed text to the output file.
        
        Args:
            processed_text (str): The fully corrected transcript.
        """
        with open(self.output_path, 'w', encoding='utf-8') as file:
            file.write(processed_text)
        print(f"Results written to {self.output_path}")
