pip install openai pymupdf pandas pydataframer tenacity pyyaml requests

FILES_DIR = Path("files")

sample_dirs = sorted([d for d in FILES_DIR.iterdir() if d.is_dir()])

samples = []
for sample_dir in sample_dirs:
    pdf_paths = list(sample_dir.glob("*.pdf"))
    qa_path = sample_dir / "qa-pairs.csv"

    pdf_path = pdf_paths[0] if pdf_paths else None
    qa_df = pd.read_csv(qa_path) if qa_path.exists() else None

    samples.append({
        "dir": sample_dir,
        "pdf_path": pdf_path,
        "qa_path": qa_path,
        "qa_df": qa_df,
    })

    pdf_name = pdf_path.name if pdf_path else "(no PDF)"
    q_count = len(qa_df) if qa_df is not None else 0
    print(f"  {sample_dir.name}: {pdf_name} | {q_count} Q&A pairs")

zip_buffer = io.BytesIO()
with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
    for sample in samples:
        for file_path in sorted(sample["dir"].iterdir()):
            arcname = f"{sample['dir'].name}/{file_path.name}"
            zf.write(file_path, arcname=arcname)
zip_buffer.seek(0)

dataset = df_client.dataframer.seed_datasets.create_from_zip(
    name=f"bank_statements_golden_labels_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
    description="Bank financial statement PDFs with golden Q&A pairs — multi-folder seed dataset",
    zip_file=zip_buffer,
)

dataset_id = dataset.id
print(f"Dataset ID   : {dataset_id}")
print(f"Dataset type : {dataset.dataset_type}")
print(f"Files        : {dataset.file_count}")
print(f"Folders      : {dataset.folder_count}")

spec = df_client.dataframer.specs.create(
    dataset_id=dataset_id,
    name=spec_name,
    spec_generation_model_name="anthropic/claude-opus-4-6-thinking",
    generation_objectives=(
        "Do include Bank Profile as a data property. "
        "As an example, a value could be 'CRE-heavy community bank with Total CRE >300% of Tier 1'"
    ),
    extrapolate_values=False,
    generate_distributions=True,
)

TARGET_BANK_PROFILE = "CRE-heavy community bank with Total CRE >300% of Tier 1"

for prop in updated_spec_data["data_property_variations"]:
    if prop["property_name"] == "Bank Profile":
        prop["base_distributions"] = {
            v: (100 if v == TARGET_BANK_PROFILE else 0)
            for v in prop["property_values"]
        }

CONSISTENCY_REQUIREMENT = """
Very important: the document must be 100% internally consistent. All generated numbers must be consistent with each other. You MUST use the calculator tool to verify every single numerical relationship in the document for every number.

Before using the calculator, first plan the complete structure of the document: every table, every row label, every column, every entity/location. Decide what line items exist and how they roll up into subtotals and totals. Only after that create SINGLE, COMPLETE script with ALL calculations included in that one script, using the calculator to verify/derive or compute every single number that will appear in the document. Specifically:
- This applies across ALL columns and ALL rows of ALL tables, everything must be verified in ONE script.
- You are not allowed to rely solely on mental arithmetic for any calculation.
- If the document contains multiple locations, entities, departments, or time periods, the calculator must compute numbers for ALL of them.
- Every subtotal and total must be verified as the exact sum of its component line items. Every percentage must be verified as the correct division, rounded to the displayed precision.

There is NO tolerance for even small errors — if any sum or percentage deviates from its correct value even by $0.01, this is a grave problem.

Additionally, every answer in qa-pairs.csv MUST be exactly derivable from the numbers present in the accompanying PDF.
"""

existing_requirements = updated_spec_data.get("requirements") or ""
updated_spec_data["requirements"] = existing_requirements + CONSISTENCY_REQUIREMENT

updated_yaml = yaml.dump({"spec": updated_spec_data}, allow_unicode=True, sort_keys=False)
updated_spec = df_client.dataframer.specs.update(spec_id=spec_id, content_yaml=updated_yaml)

run = df_client.dataframer.runs.create(
    spec_id=updated_spec.id,
    number_of_samples=3,
    generation_model="anthropic/claude-opus-4-6-thinking",
    revision_model="anthropic/claude-opus-4-6-thinking",
    revision_types=["conformance"],
    max_revision_cycles=1,
    filtering_types=["conformance"],
    generation_thinking_budget=16000,
    tools=["calculator"],
    skip_outline=False,
    unified_multifield=False,
)

with open(first_sample_dir / "folder.metadata") as f:
    folder_meta = json.load(f)

for key, value in folder_meta.get("generation_tags", {}).items():
    print(f"  {key}: {value}")

SYSTEM_PROMPT = (
    "You are a precise financial data extraction assistant. "
    "The user will provide the full text of a bank financial statement and ask a specific question. "
    "Answer with only the single requested value — no explanation, no units label, no extra text. "
    "If answering with a number, answer with 2 decimal places."
)

doc = fitz.open(first_pdf)
pdf_text = "\n\n".join(page.get_text() for page in doc)

eval_results = []
for _, row in qa_df.iterrows():
    question = row["question"]
    golden = str(row["answer"])
    response = openai_client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": f"Financial statement:\n\n{pdf_text}\n\nQuestion: {question}"},
        ],
        temperature=0,
    )
    model_answer = response.choices[0].message.content.strip()
    match = model_answer == golden
    eval_results.append({"question": question, "model_answer": model_answer, "golden": golden, "exact_match": match})

results_df = pd.DataFrame(eval_results)
accuracy = results_df["exact_match"].mean()
print(f"Exact-match accuracy: {accuracy:.1%} ({int(results_df['exact_match'].sum())}/{len(results_df)})")