OCR-Viewer (#1)

Co-authored-by: nathan <nathan.kehler@gmail.com>
Reviewed-on: #1

GoogleDocumentOCR/main.py

@@ -0,0 +1,149 @@
+import os
+import shutil
+from tqdm import tqdm
+from dotenv import load_dotenv
+
+from db import init_db
+from ocr import ocr_file
+from embed_and_store import process_pages
+
+load_dotenv()
+
+DOCUMENTS_DIR = "documents"
+PROCESSED_DIR = "processed"
+
+SUPPORTED_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".tif", ".gif", ".bmp", ".webp"}
+
+# ── Tag your documents here ──────────────────────────────────────────────────
+# Maps a filename substring → (nationality, corpus)
+# First matching rule wins. Falls back to ("unknown", "unknown").
+DOCUMENT_TAGS = {
+    # German records
+    "german":    ("german",   "german_records"),
+    "deutsch":   ("german",   "german_records"),
+    "Wehrmacht": ("german",   "german_records"),
+    # Canadian war diaries
+    "Calgary":   ("canadian", "war_diary_narrative"),
+    "Blackwatch": ("canadian","war_diary_narrative"),
+    "5CIB":      ("canadian", "war_diary_narrative"),
+}
+
+def tag_document(filename: str) -> tuple[str, str]:
+    for key, (nationality, corpus) in DOCUMENT_TAGS.items():
+        if key.lower() in filename.lower():
+            return nationality, corpus
+    return "unknown", "unknown"
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def main():
+    os.makedirs(DOCUMENTS_DIR, exist_ok=True)
+    os.makedirs(PROCESSED_DIR, exist_ok=True)
+    init_db()
+
+    files = [
+        f for f in os.listdir(DOCUMENTS_DIR)
+        if os.path.splitext(f)[1].lower() in SUPPORTED_EXTENSIONS
+    ]
+
+    if not files:
+        print(f"No supported files found in '{DOCUMENTS_DIR}'. Drop PDFs or images there and re-run.")
+        return
+
+    print(f"Found {len(files)} file(s) to process.\n")
+
+    for filename in tqdm(files, desc="Processing documents"):
+        file_path = os.path.join(DOCUMENTS_DIR, filename)
+        nationality, corpus = tag_document(filename)
+        print(f"\n--- {filename} --- [{nationality} / {corpus}]")
+
+        try:
+            pages = ocr_file(file_path)
+            print(f"  OCR complete: {len(pages)} page(s)")
+
+            process_pages(filename, pages, nationality=nationality, corpus=corpus)
+
+            shutil.move(file_path, os.path.join(PROCESSED_DIR, filename))
+            print(f"  Moved to {PROCESSED_DIR}/")
+
+        except Exception as e:
+            print(f"  ERROR processing {filename}: {e}")
+
+    print("\nDone. All files processed.")
+
+
+def search(query: str, top_k: int = 5, nationality: str = None, corpus: str = None):
+    from embed_and_store import embed_texts
+    from db import semantic_search
+
+    filters = []
+    if nationality: filters.append(f"nationality={nationality}")
+    if corpus:      filters.append(f"corpus={corpus}")
+    filter_str = f" [{', '.join(filters)}]" if filters else ""
+    print(f"\nSearching for: '{query}'{filter_str}\n")
+
+    embedding = embed_texts([query])[0]
+    results = semantic_search(embedding, top_k=top_k, nationality=nationality, corpus=corpus)
+
+    for rank, (filename, page_num, chunk_index, text, nat, corp, similarity) in enumerate(results, 1):
+        print(f"[{rank}] {filename} — page {page_num}, chunk {chunk_index} [{nat} / {corp}] (similarity: {similarity:.3f})")
+        print(f"    {text[:200]}...")
+        print()
+def ask(question: str, top_k: int = 5, nationality: str = None, corpus: str = None):
+    from embed_and_store import embed_texts
+    from db import semantic_search
+    import os
+    from openai import OpenAI
+
+    filters = []
+    if nationality: filters.append(f"nationality={nationality}")
+    if corpus:      filters.append(f"corpus={corpus}")
+    filter_str = f" [{', '.join(filters)}]" if filters else ""
+    print(f"\nAsking: '{question}'{filter_str}\n")
+
+    # Step 1 - find relevant chunks
+    embedding = embed_texts([question])[0]
+    results = semantic_search(embedding, top_k=top_k, nationality=nationality, corpus=corpus)
+
+    if not results:
+        print("No relevant chunks found.")
+        return
+
+    # Step 2 - build context from chunks
+    context = ""
+    for filename, page_num, chunk_index, text, nat, corp, similarity in results:
+        context += f"[Source: {filename}, page {page_num}]\n{text}\n\n"
+
+    # Step 3 - send to LLM
+    client = OpenAI(
+        api_key=os.getenv("DEEPINFRA_API_KEY"),
+        base_url="https://api.deepinfra.com/v1/openai",
+    )
+
+    prompt = f"""You are a WWII military historian. Answer the question below using only the provided source documents.
+Cite the source file and page number for each claim you make.
+
+QUESTION: {question}
+
+SOURCE DOCUMENTS:
+{context}
+
+Answer:"""
+
+    response = client.chat.completions.create(
+        model="meta-llama/Meta-Llama-3.1-8B-Instruct",
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=1000,
+    )
+
+    print("\nQUESTION:", question)
+    print("\nANSWER:")
+    print(response.choices[0].message.content)
+    print("\nSOURCES USED:")
+    for filename, page_num, chunk_index, text, nat, corp, similarity in results:
+        print(f"  - {filename} page {page_num} [{nat} / {corp}] (similarity: {similarity:.3f})")
+
+if __name__ == "__main__":
+    # main()
+    # search("is this document in german?", nationality="german")
+    ask("What was the Calgary Highlanders doing in October 1944?", nationality="canadian", corpus="war_diary_narrative")