청크 간 맥락 유지하기

import tiktoken

enc = tiktoken.get_encoding('cl100k_base')

def chunk_with_overlap(text, max_tokens=1000, overlap=200):
    tokens = enc.encode(text)
    chunks = []
    step = max_tokens - overlap
    i = 0
    while i < len(tokens):
        chunk = tokens[i:i + max_tokens]
        chunks.append(enc.decode(chunk))
        i += step
    return chunks

chunks = chunk_with_overlap(document, max_tokens=1000, overlap=200)
print(f'{len(chunks)} chunks with 200-token overlap')

def strip_overlap(chunks, overlap_tokens=200):
    '''Remove the leading overlap from each chunk (except the first).'''
    cleaned = [chunks[0]]
    for chunk in chunks[1:]:
        tokens = enc.encode(chunk)
        trimmed = enc.decode(tokens[overlap_tokens:])
        cleaned.append(trimmed)
    return cleaned

import openai
client = openai.OpenAI(api_key='sk-...')

def process_with_rolling_summary(chunks):
    rolling_summary = ''
    results = []

    for i, chunk in enumerate(chunks):
        context = ''
        if rolling_summary:
            context = f'Summary of previous sections:\n{rolling_summary}\n\n'

        prompt = context + f'Current section:\n{chunk}'
        resp = client.chat.completions.create(
            model='gpt-4o',
            messages=[
                {'role': 'system', 'content': 'Answer questions or summarize, using prior context.'},
                {'role': 'user', 'content': prompt}
            ]
        )
        result = resp.choices[0].message.content
        results.append(result)

        # Update rolling summary
        rolling_summary = update_rolling_summary(rolling_summary, chunk)

    return results

def update_rolling_summary(current_summary, new_chunk, max_words=150):
    if not current_summary:
        prompt = f'Summarize the following in under {max_words} words:\n\n{new_chunk}'
    else:
        prompt = (
            f'Current running summary (under {max_words} words):\n{current_summary}\n\n'
            f'New section to incorporate:\n{new_chunk}\n\n'
            f'Update the summary to include the new section. Stay under {max_words} words.'
        )
    resp = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}]
    )
    return resp.choices[0].message.content

def build_prompt_with_metadata(chunk, metadata):
    header = (
        f'Document: {metadata["title"]}\n'
        f'Chapter: {metadata["chapter"]}\n'
        f'Section: {metadata["section"]}\n'
        f'Page: {metadata["page"]}\n'
        f'Chunk: {metadata["chunk_index"] + 1} of {metadata["total_chunks"]}\n'
        '---\n'
    )
    return header + chunk

prompt = build_prompt_with_metadata(
    chunk=chunks[5],
    metadata={
        'title': 'Annual Report 2024',
        'chapter': '3. Financial Results',
        'section': '3.2 Revenue Breakdown',
        'page': 42,
        'chunk_index': 5,
        'total_chunks': 120
    }
)

def inject_page_markers(pages):
    '''pages: list of strings, one per page.'''
    marked = []
    for i, page_text in enumerate(pages):
        marked.append(f'[PAGE {i+1}]\n{page_text}')
    return '\n\n'.join(marked)

# When the model sees [PAGE 12] in context, it can say
# 'As defined on page 12...' in its output, enabling traceability.
document_with_markers = inject_page_markers(pdf_pages)
chunks = chunk_with_overlap(document_with_markers)

def process_document(pages, doc_metadata):
    # Step 1: inject page markers
    full_text = inject_page_markers(pages)
    # Step 2: chunk with overlap
    chunks = chunk_with_overlap(full_text, max_tokens=900, overlap=150)
    total = len(chunks)
    rolling_summary = ''
    results = []

    for i, chunk in enumerate(chunks):
        meta = {**doc_metadata, 'chunk_index': i, 'total_chunks': total}
        prompt = build_prompt_with_metadata(chunk, meta)
        if rolling_summary:
            prompt = 'Prior context:\n' + rolling_summary + '\n\n' + prompt
        result = call_llm(prompt)
        results.append(result)
        rolling_summary = update_rolling_summary(rolling_summary, chunk)

    return results

test_questions = [
    {
        'question': 'What is the definition of "net recurring revenue" used in this report?',
        'defined_in_chunk': 2,
        'asked_in_chunk': 9,
        'expected_keywords': ['net recurring revenue', 'subscription', 'exclude']
    }
]

def evaluate_context_retention(pipeline_results, test_questions):
    for test in test_questions:
        answer = pipeline_results[test['asked_in_chunk']]
        for kw in test['expected_keywords']:
            if kw.lower() not in answer.lower():
                print(f'FAIL: missing "{kw}" in answer to chunk {test["asked_in_chunk"]}')