150 lines
5.4 KiB
Python
150 lines
5.4 KiB
Python
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")
|