DocLLM — Custom Transformer for PDF Question Answering
A custom autoregressive Transformer LLM built from scratch in PyTorch, trained to read PDF documents and answer questions with page citations.
Architecture
- d_model: 256 | Layers: 6 | Heads: 8
- Context window: 256 tokens
- Parameters: 4,997,888
- Rotary Position Embeddings (RoPE)
- Multi-Head Causal Self-Attention + KV-Cache
- SwiGLU Feed-Forward Network
- RMSNorm
Usage
from inference.pdf_qa import DocLLMPipeline
from data.pdf_loader import PDFLoader
from huggingface_hub import hf_hub_download
ckpt_path = hf_hub_download("Ystar124/formtgpt", "pytorch_model.bin")
tok_path = hf_hub_download("Ystar124/formtgpt", "tokenizer.json")
pipeline = DocLLMPipeline.from_checkpoint(ckpt_path, tokenizer_path=tok_path)
doc = PDFLoader.load_pdf("your_document.pdf")
result = pipeline.ask_document(doc, "What is this document about?")
print(result["answer"])
print("Pages cited:", result["cited_pages"])