OCR-Viewer (#1)
Co-authored-by: nathan <nathan.kehler@gmail.com> Reviewed-on: #1
This commit is contained in:
108
GoogleDocumentOCR/db.py
Normal file
108
GoogleDocumentOCR/db.py
Normal file
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import psycopg2
|
||||
from psycopg2.extras import execute_values
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
def get_conn():
|
||||
return psycopg2.connect(
|
||||
host=os.getenv("DB_HOST"),
|
||||
port=os.getenv("DB_PORT"),
|
||||
dbname=os.getenv("DB_NAME"),
|
||||
user=os.getenv("DB_USER"),
|
||||
password=os.getenv("DB_PASSWORD"),
|
||||
)
|
||||
|
||||
def init_db():
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
|
||||
cur.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_chunks (
|
||||
id SERIAL PRIMARY KEY,
|
||||
filename TEXT NOT NULL,
|
||||
page_num INTEGER,
|
||||
chunk_index INTEGER,
|
||||
text TEXT NOT NULL,
|
||||
embedding vector(768),
|
||||
nationality TEXT DEFAULT 'unknown',
|
||||
corpus TEXT DEFAULT 'unknown',
|
||||
created_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
""")
|
||||
# Migrate existing tables that predate nationality/corpus columns
|
||||
for col, default in [("nationality", "unknown"), ("corpus", "unknown")]:
|
||||
cur.execute(f"""
|
||||
ALTER TABLE document_chunks
|
||||
ADD COLUMN IF NOT EXISTS {col} TEXT DEFAULT '{default}';
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS document_chunks_embedding_idx
|
||||
ON document_chunks
|
||||
USING ivfflat (embedding vector_cosine_ops)
|
||||
WITH (lists = 100);
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS document_chunks_nationality_idx
|
||||
ON document_chunks (nationality);
|
||||
""")
|
||||
cur.execute("""
|
||||
CREATE INDEX IF NOT EXISTS document_chunks_corpus_idx
|
||||
ON document_chunks (corpus);
|
||||
""")
|
||||
conn.commit()
|
||||
cur.close()
|
||||
conn.close()
|
||||
print("DB initialized.")
|
||||
|
||||
def upsert_chunks(rows: list[dict]):
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
execute_values(
|
||||
cur,
|
||||
"""
|
||||
INSERT INTO document_chunks (filename, page_num, chunk_index, text, embedding, nationality, corpus)
|
||||
VALUES %s
|
||||
""",
|
||||
[
|
||||
(
|
||||
r["filename"], r["page_num"], r["chunk_index"], r["text"], r["embedding"],
|
||||
r.get("nationality", "unknown"), r.get("corpus", "unknown"),
|
||||
)
|
||||
for r in rows
|
||||
],
|
||||
template="(%s, %s, %s, %s, %s::vector, %s, %s)"
|
||||
)
|
||||
conn.commit()
|
||||
cur.close()
|
||||
conn.close()
|
||||
|
||||
def semantic_search(query_embedding: list[float], top_k: int = 5,
|
||||
nationality: str = None, corpus: str = None):
|
||||
conn = get_conn()
|
||||
cur = conn.cursor()
|
||||
filters = []
|
||||
params = []
|
||||
if nationality:
|
||||
filters.append("nationality = %s")
|
||||
params.append(nationality)
|
||||
if corpus:
|
||||
filters.append("corpus = %s")
|
||||
params.append(corpus)
|
||||
where = ("WHERE " + " AND ".join(filters)) if filters else ""
|
||||
cur.execute(
|
||||
f"""
|
||||
SELECT filename, page_num, chunk_index, text, nationality, corpus,
|
||||
1 - (embedding <=> %s::vector) AS similarity
|
||||
FROM document_chunks
|
||||
{where}
|
||||
ORDER BY embedding <=> %s::vector
|
||||
LIMIT %s;
|
||||
""",
|
||||
[query_embedding] + params + [query_embedding, top_k]
|
||||
)
|
||||
rows = cur.fetchall()
|
||||
cur.close()
|
||||
conn.close()
|
||||
return rows
|
||||
73
GoogleDocumentOCR/embed_and_store.py
Normal file
73
GoogleDocumentOCR/embed_and_store.py
Normal file
@@ -0,0 +1,73 @@
|
||||
import os
|
||||
from openai import OpenAI
|
||||
from db import upsert_chunks
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("DEEPINFRA_API_KEY"),
|
||||
base_url="https://api.deepinfra.com/v1/openai",
|
||||
)
|
||||
|
||||
CHUNK_SIZE = 500
|
||||
CHUNK_OVERLAP = 50
|
||||
EMBEDDING_MODEL = "BAAI/bge-base-en-v1.5"
|
||||
EMBEDDING_DIMS = 768 # bge-base produces 768-dim vectors
|
||||
|
||||
|
||||
def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
|
||||
words = text.split()
|
||||
chunks = []
|
||||
start = 0
|
||||
while start < len(words):
|
||||
end = start + chunk_size
|
||||
chunks.append(" ".join(words[start:end]))
|
||||
start += chunk_size - overlap
|
||||
return [c for c in chunks if c.strip()]
|
||||
|
||||
|
||||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
all_embeddings = []
|
||||
batch_size = 100
|
||||
for i in range(0, len(texts), batch_size):
|
||||
batch = texts[i:i + batch_size]
|
||||
response = client.embeddings.create(
|
||||
input=batch,
|
||||
model=EMBEDDING_MODEL,
|
||||
)
|
||||
all_embeddings.extend([item.embedding for item in response.data])
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def process_pages(filename: str, pages: list[dict],
|
||||
nationality: str = "unknown", corpus: str = "unknown"):
|
||||
rows = []
|
||||
all_chunks = []
|
||||
|
||||
for page in pages:
|
||||
chunks = chunk_text(page["text"])
|
||||
for idx, chunk in enumerate(chunks):
|
||||
all_chunks.append({
|
||||
"filename": filename,
|
||||
"page_num": page["page_num"],
|
||||
"chunk_index": idx,
|
||||
"text": chunk,
|
||||
"nationality": nationality,
|
||||
"corpus": corpus,
|
||||
})
|
||||
|
||||
if not all_chunks:
|
||||
print(f" No text extracted from {filename}, skipping.")
|
||||
return
|
||||
|
||||
print(f" Embedding {len(all_chunks)} chunks from {filename}...")
|
||||
texts = [c["text"] for c in all_chunks]
|
||||
embeddings = embed_texts(texts)
|
||||
|
||||
for chunk, embedding in zip(all_chunks, embeddings):
|
||||
chunk["embedding"] = embedding
|
||||
rows.append(chunk)
|
||||
|
||||
upsert_chunks(rows)
|
||||
print(f" Stored {len(rows)} chunks for {filename}.")
|
||||
149
GoogleDocumentOCR/main.py
Normal file
149
GoogleDocumentOCR/main.py
Normal file
@@ -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")
|
||||
94
GoogleDocumentOCR/ocr.py
Normal file
94
GoogleDocumentOCR/ocr.py
Normal file
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
import io
|
||||
from google.api_core.client_options import ClientOptions
|
||||
from google.cloud import documentai
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
PROJECT_ID = os.getenv("GOOGLE_PROJECT_ID")
|
||||
LOCATION = os.getenv("GOOGLE_LOCATION")
|
||||
PROCESSOR_ID = os.getenv("GOOGLE_PROCESSOR_ID")
|
||||
|
||||
PAGE_LIMIT = 14 # stay under the 30 page limit
|
||||
|
||||
|
||||
def get_client():
|
||||
opts = ClientOptions(api_endpoint=f"{LOCATION}-documentai.googleapis.com")
|
||||
return documentai.DocumentProcessorServiceClient(client_options=opts)
|
||||
|
||||
|
||||
def split_pdf(content: bytes, page_limit: int = PAGE_LIMIT) -> list[bytes]:
|
||||
"""Split PDF bytes into chunks of page_limit pages."""
|
||||
import pypdf
|
||||
reader = pypdf.PdfReader(io.BytesIO(content))
|
||||
total_pages = len(reader.pages)
|
||||
chunks = []
|
||||
for start in range(0, total_pages, page_limit):
|
||||
writer = pypdf.PdfWriter()
|
||||
for i in range(start, min(start + page_limit, total_pages)):
|
||||
writer.add_page(reader.pages[i])
|
||||
buf = io.BytesIO()
|
||||
writer.write(buf)
|
||||
chunks.append(buf.getvalue())
|
||||
return chunks
|
||||
|
||||
|
||||
def ocr_chunk(client, name: str, content: bytes, mime_type: str, page_offset: int) -> list[dict]:
|
||||
"""OCR a single chunk and return pages with corrected page numbers."""
|
||||
raw_doc = documentai.RawDocument(content=content, mime_type=mime_type)
|
||||
request = documentai.ProcessRequest(name=name, raw_document=raw_doc)
|
||||
result = client.process_document(request=request)
|
||||
doc = result.document
|
||||
|
||||
pages = []
|
||||
for i, page in enumerate(doc.pages):
|
||||
page_text_parts = []
|
||||
for block in page.blocks:
|
||||
seg = block.layout.text_anchor.text_segments
|
||||
for s in seg:
|
||||
start = int(s.start_index) if s.start_index else 0
|
||||
end = int(s.end_index)
|
||||
page_text_parts.append(doc.text[start:end])
|
||||
page_text = "\n".join(page_text_parts).strip()
|
||||
if page_text:
|
||||
pages.append({"page_num": page_offset + i + 1, "text": page_text})
|
||||
return pages
|
||||
|
||||
|
||||
def ocr_file(file_path: str) -> list[dict]:
|
||||
client = get_client()
|
||||
name = client.processor_path(PROJECT_ID, LOCATION, PROCESSOR_ID)
|
||||
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
mime_map = {
|
||||
".pdf": "application/pdf",
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".tiff": "image/tiff",
|
||||
".tif": "image/tiff",
|
||||
".gif": "image/gif",
|
||||
".bmp": "image/bmp",
|
||||
".webp": "image/webp",
|
||||
}
|
||||
mime_type = mime_map.get(ext, "application/pdf")
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
|
||||
# Only PDFs need splitting
|
||||
if mime_type == "application/pdf":
|
||||
chunks = split_pdf(content)
|
||||
print(f" Split into {len(chunks)} chunk(s) of up to {PAGE_LIMIT} pages")
|
||||
else:
|
||||
chunks = [content]
|
||||
|
||||
all_pages = []
|
||||
for chunk_idx, chunk in enumerate(chunks):
|
||||
page_offset = chunk_idx * PAGE_LIMIT
|
||||
print(f" OCR chunk {chunk_idx + 1}/{len(chunks)}...")
|
||||
pages = ocr_chunk(client, name, chunk, mime_type, page_offset)
|
||||
all_pages.extend(pages)
|
||||
|
||||
return all_pages
|
||||
473
GoogleDocumentOCR/web_app.py
Normal file
473
GoogleDocumentOCR/web_app.py
Normal file
@@ -0,0 +1,473 @@
|
||||
"""
|
||||
web_app.py — Terminal-style search UI
|
||||
Run: python GoogleDocumentOCR/web_app.py
|
||||
Then open: http://localhost:5000
|
||||
"""
|
||||
|
||||
import os, sys
|
||||
from pathlib import Path
|
||||
from flask import Flask, request, jsonify, render_template_string
|
||||
from dotenv import load_dotenv
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
load_dotenv()
|
||||
|
||||
app = Flask(__name__)
|
||||
|
||||
# ── HTML template ──────────────────────────────────────────────────────────────
|
||||
HTML = """<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>WAR DIARY SEARCH // HISTORICAL RECORDS SYSTEM</title>
|
||||
<style>
|
||||
:root {
|
||||
--green: #00ff41;
|
||||
--dimgreen: #00a82b;
|
||||
--amber: #ffb000;
|
||||
--red: #ff4444;
|
||||
--bg: #0a0a0a;
|
||||
--bg2: #0f0f0f;
|
||||
--border: #1a3a1a;
|
||||
}
|
||||
|
||||
* { box-sizing: border-box; margin: 0; padding: 0; }
|
||||
|
||||
body {
|
||||
background: var(--bg);
|
||||
color: var(--green);
|
||||
font-family: 'Courier New', Courier, monospace;
|
||||
font-size: 14px;
|
||||
height: 100vh;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
/* ── Header ── */
|
||||
#header {
|
||||
border-bottom: 1px solid var(--border);
|
||||
padding: 10px 16px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
flex-shrink: 0;
|
||||
background: var(--bg2);
|
||||
}
|
||||
#header .title {
|
||||
color: var(--amber);
|
||||
font-size: 13px;
|
||||
letter-spacing: 2px;
|
||||
text-transform: uppercase;
|
||||
}
|
||||
#header .status {
|
||||
font-size: 11px;
|
||||
color: var(--dimgreen);
|
||||
}
|
||||
#header .status span { color: var(--green); }
|
||||
|
||||
/* ── Filters bar ── */
|
||||
#filters {
|
||||
padding: 6px 16px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
display: flex;
|
||||
gap: 20px;
|
||||
align-items: center;
|
||||
font-size: 12px;
|
||||
color: var(--dimgreen);
|
||||
flex-shrink: 0;
|
||||
background: var(--bg2);
|
||||
}
|
||||
#filters label { color: var(--dimgreen); }
|
||||
#filters select {
|
||||
background: #000;
|
||||
color: var(--green);
|
||||
border: 1px solid var(--border);
|
||||
font-family: inherit;
|
||||
font-size: 12px;
|
||||
padding: 2px 6px;
|
||||
outline: none;
|
||||
cursor: pointer;
|
||||
}
|
||||
#filters select:focus { border-color: var(--green); }
|
||||
|
||||
/* ── Chat window ── */
|
||||
#chat {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
padding: 16px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 18px;
|
||||
scroll-behavior: smooth;
|
||||
}
|
||||
#chat::-webkit-scrollbar { width: 6px; }
|
||||
#chat::-webkit-scrollbar-track { background: #000; }
|
||||
#chat::-webkit-scrollbar-thumb { background: var(--border); }
|
||||
|
||||
/* ── Boot message ── */
|
||||
.boot {
|
||||
color: var(--dimgreen);
|
||||
font-size: 12px;
|
||||
line-height: 1.8;
|
||||
border-bottom: 1px solid var(--border);
|
||||
padding-bottom: 14px;
|
||||
}
|
||||
.boot .hi { color: var(--amber); }
|
||||
|
||||
/* ── Exchange (Q+A pair) ── */
|
||||
.exchange {}
|
||||
|
||||
.query-line {
|
||||
color: var(--amber);
|
||||
margin-bottom: 8px;
|
||||
word-break: break-word;
|
||||
}
|
||||
.query-line::before { content: '> '; color: var(--dimgreen); }
|
||||
|
||||
.answer-block {
|
||||
color: var(--green);
|
||||
line-height: 1.7;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
padding-left: 14px;
|
||||
border-left: 2px solid var(--border);
|
||||
}
|
||||
|
||||
.sources-block {
|
||||
margin-top: 10px;
|
||||
padding-left: 14px;
|
||||
border-left: 2px solid var(--border);
|
||||
}
|
||||
.sources-label {
|
||||
color: var(--dimgreen);
|
||||
font-size: 11px;
|
||||
letter-spacing: 1px;
|
||||
margin-bottom: 4px;
|
||||
}
|
||||
.source-row {
|
||||
font-size: 11px;
|
||||
color: #336633;
|
||||
line-height: 1.6;
|
||||
}
|
||||
.source-row .sim { color: var(--dimgreen); }
|
||||
|
||||
/* ── Thinking indicator ── */
|
||||
.thinking {
|
||||
color: var(--dimgreen);
|
||||
font-size: 12px;
|
||||
}
|
||||
.dot-anim::after {
|
||||
content: '';
|
||||
animation: dots 1.2s steps(4, end) infinite;
|
||||
}
|
||||
@keyframes dots {
|
||||
0% { content: ''; }
|
||||
25% { content: '.'; }
|
||||
50% { content: '..'; }
|
||||
75% { content: '...'; }
|
||||
100% { content: ''; }
|
||||
}
|
||||
|
||||
.error-line { color: var(--red); padding-left: 14px; }
|
||||
|
||||
/* ── Input bar ── */
|
||||
#inputbar {
|
||||
border-top: 1px solid var(--border);
|
||||
padding: 12px 16px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
flex-shrink: 0;
|
||||
background: var(--bg2);
|
||||
}
|
||||
#prompt-symbol {
|
||||
color: var(--amber);
|
||||
font-size: 15px;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
#question {
|
||||
flex: 1;
|
||||
background: transparent;
|
||||
border: none;
|
||||
outline: none;
|
||||
color: var(--green);
|
||||
font-family: inherit;
|
||||
font-size: 14px;
|
||||
caret-color: var(--green);
|
||||
}
|
||||
#question::placeholder { color: #1a3a1a; }
|
||||
|
||||
#send-btn {
|
||||
background: transparent;
|
||||
border: 1px solid var(--border);
|
||||
color: var(--dimgreen);
|
||||
font-family: inherit;
|
||||
font-size: 12px;
|
||||
padding: 4px 10px;
|
||||
cursor: pointer;
|
||||
letter-spacing: 1px;
|
||||
transition: color 0.15s, border-color 0.15s;
|
||||
}
|
||||
#send-btn:hover:not(:disabled) {
|
||||
color: var(--green);
|
||||
border-color: var(--green);
|
||||
}
|
||||
#send-btn:disabled { opacity: 0.3; cursor: not-allowed; }
|
||||
|
||||
/* ── Scanline overlay ── */
|
||||
body::after {
|
||||
content: '';
|
||||
position: fixed;
|
||||
inset: 0;
|
||||
background: repeating-linear-gradient(
|
||||
0deg,
|
||||
transparent,
|
||||
transparent 2px,
|
||||
rgba(0,0,0,0.07) 2px,
|
||||
rgba(0,0,0,0.07) 4px
|
||||
);
|
||||
pointer-events: none;
|
||||
z-index: 999;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<div id="header">
|
||||
<div class="title">■ WAR DIARY SEARCH // HISTORICAL RECORDS SYSTEM</div>
|
||||
<div class="status">DB STATUS: <span>ONLINE</span> | CORPUS: <span>CANADIAN WWII</span></div>
|
||||
</div>
|
||||
|
||||
<div id="filters">
|
||||
<span>FILTERS:</span>
|
||||
<label>NATIONALITY
|
||||
<select id="nationality">
|
||||
<option value="">ALL</option>
|
||||
<option value="canadian" selected>CANADIAN</option>
|
||||
<option value="german">GERMAN</option>
|
||||
<option value="unknown">UNKNOWN</option>
|
||||
</select>
|
||||
</label>
|
||||
<label>CORPUS
|
||||
<select id="corpus">
|
||||
<option value="">ALL</option>
|
||||
<option value="war_diary_narrative" selected>WAR DIARY NARRATIVE</option>
|
||||
<option value="war_diary_appendix">WAR DIARY APPENDIX</option>
|
||||
<option value="german_records">GERMAN RECORDS</option>
|
||||
</select>
|
||||
</label>
|
||||
<label>RESULTS
|
||||
<select id="top_k">
|
||||
<option value="3">3</option>
|
||||
<option value="5" selected>5</option>
|
||||
<option value="10">10</option>
|
||||
</select>
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<div id="chat">
|
||||
<div class="boot">
|
||||
<div class="hi">HISTORICAL RECORDS RETRIEVAL SYSTEM v1.0</div>
|
||||
<div>Corpus: Calgary Highlanders · 5 CIB · RHC Black Watch · Sep–Oct 1944</div>
|
||||
<div>Embedding model: BAAI/bge-base-en-v1.5 | LLM: Meta-Llama-3.1-8B</div>
|
||||
<div> </div>
|
||||
<div>Type a question and press ENTER or [SEND].</div>
|
||||
<div>Example: <span style="color:var(--green)">What was the Calgary Highlanders doing on October 23, 1944?</span></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="inputbar">
|
||||
<span id="prompt-symbol">></span>
|
||||
<input id="question" type="text"
|
||||
placeholder="Ask about the war diaries..."
|
||||
autocomplete="off" spellcheck="false" autofocus />
|
||||
<button id="send-btn">SEND</button>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const chatEl = document.getElementById('chat');
|
||||
const inputEl = document.getElementById('question');
|
||||
const sendBtn = document.getElementById('send-btn');
|
||||
|
||||
function escHtml(s) {
|
||||
return s.replace(/&/g,'&').replace(/</g,'<').replace(/>/g,'>');
|
||||
}
|
||||
|
||||
function scrollBottom() {
|
||||
chatEl.scrollTop = chatEl.scrollHeight;
|
||||
}
|
||||
|
||||
function addExchange(question, answer, sources, isError) {
|
||||
const div = document.createElement('div');
|
||||
div.className = 'exchange';
|
||||
|
||||
const qLine = document.createElement('div');
|
||||
qLine.className = 'query-line';
|
||||
qLine.textContent = question;
|
||||
div.appendChild(qLine);
|
||||
|
||||
if (isError) {
|
||||
const err = document.createElement('div');
|
||||
err.className = 'error-line';
|
||||
err.textContent = 'ERROR: ' + answer;
|
||||
div.appendChild(err);
|
||||
} else {
|
||||
const ans = document.createElement('div');
|
||||
ans.className = 'answer-block';
|
||||
ans.textContent = answer;
|
||||
div.appendChild(ans);
|
||||
|
||||
if (sources && sources.length) {
|
||||
const sb = document.createElement('div');
|
||||
sb.className = 'sources-block';
|
||||
const lbl = document.createElement('div');
|
||||
lbl.className = 'sources-label';
|
||||
lbl.textContent = '── SOURCES (' + sources.length + ') ──────────────';
|
||||
sb.appendChild(lbl);
|
||||
sources.forEach((s, i) => {
|
||||
const row = document.createElement('div');
|
||||
row.className = 'source-row';
|
||||
const simPct = (s.similarity * 100).toFixed(1);
|
||||
row.innerHTML =
|
||||
'[' + (i+1) + '] ' + escHtml(s.filename) +
|
||||
' pg.' + s.page_num +
|
||||
' <span class="sim">sim:' + simPct + '%</span>' +
|
||||
(s.corpus ? ' [' + escHtml(s.corpus) + ']' : '');
|
||||
sb.appendChild(row);
|
||||
});
|
||||
div.appendChild(sb);
|
||||
}
|
||||
}
|
||||
|
||||
chatEl.appendChild(div);
|
||||
scrollBottom();
|
||||
}
|
||||
|
||||
async function ask() {
|
||||
const question = inputEl.value.trim();
|
||||
if (!question) return;
|
||||
|
||||
const nationality = document.getElementById('nationality').value;
|
||||
const corpus = document.getElementById('corpus').value;
|
||||
const top_k = parseInt(document.getElementById('top_k').value);
|
||||
|
||||
inputEl.value = '';
|
||||
sendBtn.disabled = true;
|
||||
inputEl.disabled = true;
|
||||
|
||||
// Thinking indicator
|
||||
const thinking = document.createElement('div');
|
||||
thinking.className = 'thinking dot-anim';
|
||||
thinking.textContent = 'SEARCHING';
|
||||
chatEl.appendChild(thinking);
|
||||
scrollBottom();
|
||||
|
||||
try {
|
||||
const res = await fetch('/ask', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ question, nationality, corpus, top_k }),
|
||||
});
|
||||
const data = await res.json();
|
||||
chatEl.removeChild(thinking);
|
||||
|
||||
if (data.error) {
|
||||
addExchange(question, data.error, null, true);
|
||||
} else {
|
||||
addExchange(question, data.answer, data.sources, false);
|
||||
}
|
||||
} catch (e) {
|
||||
chatEl.removeChild(thinking);
|
||||
addExchange(question, e.message, null, true);
|
||||
}
|
||||
|
||||
sendBtn.disabled = false;
|
||||
inputEl.disabled = false;
|
||||
inputEl.focus();
|
||||
}
|
||||
|
||||
sendBtn.addEventListener('click', ask);
|
||||
inputEl.addEventListener('keydown', e => { if (e.key === 'Enter') ask(); });
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
|
||||
# ── API endpoint ───────────────────────────────────────────────────────────────
|
||||
@app.route("/")
|
||||
def index():
|
||||
return render_template_string(HTML)
|
||||
|
||||
|
||||
@app.route("/ask", methods=["POST"])
|
||||
def ask():
|
||||
from embed_and_store import embed_texts
|
||||
from db import semantic_search
|
||||
from openai import OpenAI
|
||||
|
||||
data = request.get_json(force=True)
|
||||
question = (data.get("question") or "").strip()
|
||||
nationality = data.get("nationality") or None
|
||||
corpus = data.get("corpus") or None
|
||||
top_k = int(data.get("top_k", 5))
|
||||
|
||||
if not question:
|
||||
return jsonify({"error": "No question provided."}), 400
|
||||
|
||||
try:
|
||||
# 1. Embed + search
|
||||
embedding = embed_texts([question])[0]
|
||||
rows = semantic_search(embedding, top_k=top_k,
|
||||
nationality=nationality, corpus=corpus)
|
||||
|
||||
if not rows:
|
||||
return jsonify({
|
||||
"answer": "No relevant records found for that query with the current filters.",
|
||||
"sources": []
|
||||
})
|
||||
|
||||
# 2. Build context
|
||||
context = ""
|
||||
sources = []
|
||||
for filename, page_num, chunk_index, text, nat, corp, similarity in rows:
|
||||
context += f"[Source: {filename}, page {page_num}]\n{text}\n\n"
|
||||
sources.append({
|
||||
"filename": filename,
|
||||
"page_num": page_num,
|
||||
"similarity": round(similarity, 4),
|
||||
"corpus": corp,
|
||||
})
|
||||
|
||||
# 3. LLM answer
|
||||
client = OpenAI(
|
||||
api_key=os.getenv("DEEPINFRA_API_KEY"),
|
||||
base_url="https://api.deepinfra.com/v1/openai",
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"You are a WWII military historian. Answer the question using only "
|
||||
"the provided source documents. Cite source file and page number "
|
||||
"for each claim.\n\n"
|
||||
f"QUESTION: {question}\n\n"
|
||||
f"SOURCE DOCUMENTS:\n{context}\n\nAnswer:"
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
max_tokens=1200,
|
||||
)
|
||||
|
||||
answer = response.choices[0].message.content.strip()
|
||||
return jsonify({"answer": answer, "sources": sources})
|
||||
|
||||
except Exception as e:
|
||||
return jsonify({"error": str(e)}), 500
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("\n WAR DIARY SEARCH — http://localhost:5000\n")
|
||||
app.run(debug=False, port=5000)
|
||||
|
||||
Reference in New Issue
Block a user