mistocr 0.0.4__tar.gz → 0.1.5__tar.gz

{mistocr-0.0.4/mistocr.egg-info → mistocr-0.1.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mistocr
-Version: 0.0.4
+Version: 0.1.5
 Summary: Simple batch OCR for PDFs using Mistral's state-of-the-art vision model
 Home-page: https://github.com/franckalbinet/mistocr
 Author: Solveit
@@ -22,6 +22,7 @@ Requires-Dist: fastcore
 Requires-Dist: mistralai
 Requires-Dist: pillow
 Requires-Dist: dotenv
+Requires-Dist: lisette
 Provides-Extra: dev
 Dynamic: author
 Dynamic: author-email

mistocr-0.1.5/mistocr/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.5"

{mistocr-0.0.4 → mistocr-0.1.5}/mistocr/_modidx.py

@@ -19,4 +19,13 @@ d = { 'settings': { 'branch': 'main',
                               'mistocr.core.save_pages': ('core.html#save_pages', 'mistocr/core.py'),
                               'mistocr.core.submit_batch': ('core.html#submit_batch', 'mistocr/core.py'),
                               'mistocr.core.upload_pdf': ('core.html#upload_pdf', 'mistocr/core.py'),
-                              'mistocr.core.wait_for_job': ('core.html#wait_for_job', 'mistocr/core.py')}}}
+                              'mistocr.core.wait_for_job': ('core.html#wait_for_job', 'mistocr/core.py')},
+            'mistocr.refine': { 'mistocr.refine.HeadingCorrections': ('refine.html#headingcorrections', 'mistocr/refine.py'),
+                                'mistocr.refine.add_pg_hdgs': ('refine.html#add_pg_hdgs', 'mistocr/refine.py'),
+                                'mistocr.refine.apply_hdg_fixes': ('refine.html#apply_hdg_fixes', 'mistocr/refine.py'),
+                                'mistocr.refine.fix_hdg_hierarchy': ('refine.html#fix_hdg_hierarchy', 'mistocr/refine.py'),
+                                'mistocr.refine.fix_md_hdgs': ('refine.html#fix_md_hdgs', 'mistocr/refine.py'),
+                                'mistocr.refine.fmt_hdgs_idx': ('refine.html#fmt_hdgs_idx', 'mistocr/refine.py'),
+                                'mistocr.refine.get_hdgs': ('refine.html#get_hdgs', 'mistocr/refine.py'),
+                                'mistocr.refine.mk_fixes_lut': ('refine.html#mk_fixes_lut', 'mistocr/refine.py'),
+                                'mistocr.refine.read_pgs_pg': ('refine.html#read_pgs_pg', 'mistocr/refine.py')}}}

{mistocr-0.0.4 → mistocr-0.1.5}/mistocr/core.py

@@ -110,11 +110,11 @@ def save_images(
 # %% ../nbs/00_core.ipynb 32
 def save_page(
     page:dict, # Page dict, 
-    out_dir:str, # Directory to save page
+    dst:str, # Directory to save page
     img_dir:str='img' # Directory to save images
     ) -> None:
     "Save single page markdown and images"
-    (out_dir / f"page_{page['index']+1}.md").write_text(page['markdown'])
+    (dst / f"page_{page['index']+1}.md").write_text(page['markdown'])
     if page.get('images'):
         img_dir.mkdir(exist_ok=True)
         save_images(page, img_dir)
@@ -122,15 +122,15 @@ def save_page(
 # %% ../nbs/00_core.ipynb 34
 def save_pages(
     ocr_resp:dict, # OCR response, 
-    out_dir:str, # Directory to save pages, 
+    dst:str, # Directory to save pages, 
     cid:str # Custom ID
     ) -> Path: # Output directory
     "Save markdown pages and images from OCR response to output directory"
-    out_dir = Path(out_dir) / cid
-    out_dir.mkdir(parents=True, exist_ok=True)
-    img_dir = out_dir / 'img'
-    for page in ocr_resp['pages']: save_page(page, out_dir, img_dir)
-    return out_dir
+    dst = Path(dst) / cid
+    dst.mkdir(parents=True, exist_ok=True)
+    img_dir = dst / 'img'
+    for page in ocr_resp['pages']: save_page(page, dst, img_dir)
+    return dst
 
 # %% ../nbs/00_core.ipynb 40
 def _get_paths(path:str) -> list[Path]:
@@ -163,7 +163,7 @@ def _run_batch(entries:list[dict], c:Mistral, poll_interval:int=2) -> list[dict]
 # %% ../nbs/00_core.ipynb 43
 def ocr(
     path:str, # Path to PDF file or folder,
-    out_dir:str='md', # Directory to save markdown pages, 
+    dst:str='md', # Directory to save markdown pages, 
     inc_img:bool=True, # Include image in response, 
     key:str=None, # API key, 
     poll_interval:int=2 # Poll interval in seconds
@@ -172,18 +172,15 @@ def ocr(
     pdfs = _get_paths(path)
     entries, c = _prep_batch(pdfs, inc_img, key)
     results = _run_batch(entries, c, poll_interval)
-    return L([save_pages(r['response']['body'], out_dir, r['custom_id']) for r in results])
+    return L([save_pages(r['response']['body'], dst, r['custom_id']) for r in results])
 
 # %% ../nbs/00_core.ipynb 48
 def read_pgs(
     path:str, # OCR output directory, 
-    pg:int=None, # Page number
-    ) -> str:
+    join:bool=True # Join pages into single string
+    ) -> str|list[str]: # Joined string or list of page contents
     "Read specific page or all pages from OCR output directory"
     path = Path(path)
-    if pg:
-        pg_path = path / f'page_{pg}.md'
-        if not pg_path.exists(): raise ValueError(f"Page {pg} not found")
-        return pg_path.read_text()
     pgs = sorted(path.glob('page_*.md'), key=lambda p: int(p.stem.split('_')[1]))
-    return '\n\n'.join([p.read_text() for p in pgs])
+    contents = L([p.read_text() for p in pgs])
+    return '\n\n'.join(contents) if join else contents

mistocr-0.1.5/mistocr/refine.py

@@ -0,0 +1,133 @@
+"""Postprocess markdown files by fixing heading hierarchy and describint images"""
+
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_refine.ipynb.
+
+# %% auto 0
+__all__ = ['prompt_fix_hdgs', 'get_hdgs', 'add_pg_hdgs', 'read_pgs_pg', 'fmt_hdgs_idx', 'HeadingCorrections', 'fix_hdg_hierarchy',
+           'mk_fixes_lut', 'apply_hdg_fixes', 'fix_md_hdgs']
+
+# %% ../nbs/01_refine.ipynb 3
+from fastcore.all import *
+from .core import read_pgs
+from re import sub, findall, MULTILINE
+from pydantic import BaseModel
+from lisette.core import completion
+import os
+import json
+
+# %% ../nbs/01_refine.ipynb 8
+def get_hdgs(
+    md:str # Markdown file string
+    ):
+    "Return the markdown headings"
+    # Sanitize removing '#' in python snippet if any
+    md = sub(r'```[\s\S]*?```', '', md)
+    return L(findall(r'^#{1,6} .+$', md, MULTILINE))
+
+
+
+# %% ../nbs/01_refine.ipynb 9
+def add_pg_hdgs(md, n):
+    "Add page number to all headings in markdown"
+    md = sub(r'```[\s\S]*?```', '', md)
+    def repl(m): return m.group(0) + f' ... page {n}'
+    return sub(r'^#{1,6} .+$', repl, md, flags=MULTILINE)
+
+# %% ../nbs/01_refine.ipynb 11
+def read_pgs_pg(path):
+    "Read all pages of a markdown file and add page numbers to all headings"
+    pgs = read_pgs(path, join=False)
+    return L([add_pg_hdgs(pg, n) for n, pg in enumerate(pgs, 1)]).concat()
+
+# %% ../nbs/01_refine.ipynb 15
+def fmt_hdgs_idx(
+    hdgs: list[str] # List of markdown headings
+    ) -> str: # Formatted string with index
+    "Format the headings with index"
+    return '\n'.join(f"{i}. {h}" for i, h in enumerate(hdgs))
+
+
+# %% ../nbs/01_refine.ipynb 18
+class HeadingCorrections(BaseModel):
+    corrections: dict[int, str]  # index → corrected heading
+
+# %% ../nbs/01_refine.ipynb 20
+prompt_fix_hdgs = """Fix markdown heading hierarchy errors while preserving the document's intended structure.
+
+INPUT FORMAT: Each heading is prefixed with its index number (e.g., "0. # Title ... page 1")
+
+RULES - Apply these fixes in order:
+
+1. **Single H1 rule**: Documents must have exactly ONE # heading (typically the document title at the top)
+   - If index 0 is already #, then all subsequent headings (index 1+) must be ## or deeper
+   - If no H1 exists, the first major heading should be #, and all others ## or deeper
+   - NO exceptions: appendices, references, and all sections are ## or deeper after the title
+
+2. **Infer depth from numbering patterns**: If headings contain section numbers, deeper nesting means deeper heading level
+   - Parent section (e.g., "1", "2", "A") should be shallower than child (e.g., "1.1", "2.a", "A.1")
+   - Child section should be one # deeper than parent
+   - Works with any numbering: "1/1.1/1.1.1", "A/A.1/A.1.a", "I/I.A/I.A.1", etc.
+
+3. **Level jumps**: Headings can only increase by one # at a time when moving deeper
+   - Wrong: ## Section → ##### Subsection
+   - Fixed: ## Section → ### Subsection
+
+4. **Decreasing levels is OK**: Moving back up the hierarchy (### to ##) is valid for new sections
+
+OUTPUT: Return a Python dictionary mapping index to corrected heading (without the index prefix).
+IMPORTANT: Preserve the " ... page N" suffix in all corrected headings.
+Only include entries that need changes.
+
+Headings to analyze:
+{headings_list}
+"""
+
+
+# %% ../nbs/01_refine.ipynb 21
+def fix_hdg_hierarchy(
+    hdgs: list[str], # List of markdown headings
+    prompt: str=prompt_fix_hdgs, # Prompt to use
+    model: str='claude-sonnet-4-5', # Model to use
+    api_key: str=os.getenv('ANTHROPIC_API_KEY') # API key
+    ) -> dict[int, str]: # Dictionary of index → corrected heading
+    "Fix the heading hierarchy"
+    r = completion(
+        model=model, 
+        messages=[{"role": "user", "content": prompt_fix_hdgs.format(headings_list=fmt_hdgs_idx(hdgs))}], 
+        response_format=HeadingCorrections, 
+        api_key=api_key
+        )
+    return json.loads(r.choices[0].message.content)['corrections']
+
+# %% ../nbs/01_refine.ipynb 24
+def mk_fixes_lut(
+    hdgs: list[str], # List of markdown headings
+    model: str='claude-sonnet-4-5', # Model to use
+    api_key: str=os.getenv('ANTHROPIC_API_KEY') # API key
+    ) -> dict[str, str]: # Dictionary of old → new heading
+    "Make a lookup table of fixes"
+    fixes = fix_hdg_hierarchy(hdgs, model=model, api_key=api_key)
+    return {hdgs[int(k)]:v for k,v in fixes.items()}
+
+# %% ../nbs/01_refine.ipynb 27
+def apply_hdg_fixes(
+    p:str, # Page to fix
+    lut_fixes: dict[str, str], # Lookup table of fixes
+    ) -> str: # Page with fixes applied
+    "Apply the fixes to the page"
+    for old in get_hdgs(p): p = p.replace(old, lut_fixes.get(old, old))
+    return p
+
+# %% ../nbs/01_refine.ipynb 30
+def fix_md_hdgs(
+    src:str, # Source directory with markdown pages
+    model:str='claude-sonnet-4-5', # Model
+    dst:str=None, # Destination directory (None=overwrite)
+):
+    "Fix heading hierarchy in markdown document"
+    src_path,dst_path = Path(src),Path(dst) if dst else Path(src)
+    if dst_path != src_path: dst_path.mkdir(parents=True, exist_ok=True)
+    pgs_with_pg = read_pgs_pg(src_path)
+    lut = mk_fixes_lut(L([get_hdgs(pg) for pg in pgs_with_pg]).concat(), model)
+    for i,p in enumerate(pgs_with_pg, 1):
+        (dst_path/f'page_{i}.md').write_text(apply_hdg_fixes(p, lut))

{mistocr-0.0.4 → mistocr-0.1.5/mistocr.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mistocr
-Version: 0.0.4
+Version: 0.1.5
 Summary: Simple batch OCR for PDFs using Mistral's state-of-the-art vision model
 Home-page: https://github.com/franckalbinet/mistocr
 Author: Solveit
@@ -22,6 +22,7 @@ Requires-Dist: fastcore
 Requires-Dist: mistralai
 Requires-Dist: pillow
 Requires-Dist: dotenv
+Requires-Dist: lisette
 Provides-Extra: dev
 Dynamic: author
 Dynamic: author-email

{mistocr-0.0.4 → mistocr-0.1.5}/mistocr.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ setup.py
 mistocr/__init__.py
 mistocr/_modidx.py
 mistocr/core.py
+mistocr/refine.py
 mistocr.egg-info/PKG-INFO
 mistocr.egg-info/SOURCES.txt
 mistocr.egg-info/dependency_links.txt

{mistocr-0.0.4 → mistocr-0.1.5}/mistocr.egg-info/requires.txt

@@ -2,5 +2,6 @@ fastcore
 mistralai
 pillow
 dotenv
+lisette
 
 [dev]

{mistocr-0.0.4 → mistocr-0.1.5}/settings.ini

@@ -1,7 +1,7 @@
 [DEFAULT]
 repo = mistocr
 lib_name = mistocr
-version = 0.0.4
+version = 0.1.5
 min_python = 3.9
 license = apache2
 black_formatting = False
@@ -27,7 +27,7 @@ keywords = nbdev jupyter notebook python
 language = English
 status = 3
 user = franckalbinet
-requirements = fastcore mistralai pillow dotenv
+requirements = fastcore mistralai pillow dotenv lisette
 readme_nb = index.ipynb
 allowed_metadata_keys = 
 allowed_cell_metadata_keys = 

mistocr-0.0.4/mistocr/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.4"