| """
|
| FastAPI web service for document text extraction.
|
| Provides REST API endpoints for uploading and processing documents.
|
| """
|
|
|
| from fastapi import FastAPI, File, UploadFile, HTTPException, Form
|
| from fastapi.middleware.cors import CORSMiddleware
|
| from fastapi.responses import HTMLResponse, JSONResponse
|
| from fastapi.staticfiles import StaticFiles
|
| import uvicorn
|
| import tempfile
|
| import os
|
| import json
|
| from pathlib import Path
|
| from typing import List, Optional, Dict, Any
|
| import shutil
|
|
|
| from src.inference import DocumentInference
|
|
|
|
|
|
|
| app = FastAPI(
|
| title="Document Text Extraction API",
|
| description="Extract structured information from documents using Small Language Model (SLM)",
|
| version="1.0.0"
|
| )
|
|
|
|
|
| app.add_middleware(
|
| CORSMiddleware,
|
| allow_origins=["*"],
|
| allow_credentials=True,
|
| allow_methods=["*"],
|
| allow_headers=["*"],
|
| )
|
|
|
|
|
| inference_pipeline: Optional[DocumentInference] = None
|
|
|
| def get_inference_pipeline() -> DocumentInference:
|
| """Get or initialize the inference pipeline."""
|
| global inference_pipeline
|
|
|
| if inference_pipeline is None:
|
| model_path = "models/document_ner_model"
|
|
|
| if not Path(model_path).exists():
|
| raise HTTPException(
|
| status_code=503,
|
| detail="Model not found. Please train the model first by running training_pipeline.py"
|
| )
|
|
|
| try:
|
| inference_pipeline = DocumentInference(model_path)
|
| except Exception as e:
|
| raise HTTPException(
|
| status_code=503,
|
| detail=f"Failed to load model: {str(e)}"
|
| )
|
|
|
| return inference_pipeline
|
|
|
|
|
| @app.on_event("startup")
|
| async def startup_event():
|
| """Initialize the model on startup."""
|
| try:
|
| get_inference_pipeline()
|
| print("Model loaded successfully on startup")
|
| except Exception as e:
|
| print(f"Failed to load model on startup: {e}")
|
| print("Model will be loaded on first request")
|
|
|
|
|
| @app.get("/", response_class=HTMLResponse)
|
| async def root():
|
| """Serve the main HTML interface."""
|
| html_content = """
|
| <!DOCTYPE html>
|
| <html>
|
| <head>
|
| <title>Document Text Extraction</title>
|
| <style>
|
| body {
|
| font-family: Arial, sans-serif;
|
| max-width: 800px;
|
| margin: 0 auto;
|
| padding: 20px;
|
| background-color: #f5f5f5;
|
| }
|
| .container {
|
| background: white;
|
| padding: 30px;
|
| border-radius: 10px;
|
| box-shadow: 0 2px 10px rgba(0,0,0,0.1);
|
| }
|
| .header {
|
| text-align: center;
|
| color: #333;
|
| margin-bottom: 30px;
|
| }
|
| .upload-area {
|
| border: 2px dashed #ccc;
|
| padding: 40px;
|
| text-align: center;
|
| margin: 20px 0;
|
| border-radius: 8px;
|
| background-color: #fafafa;
|
| }
|
| .upload-area:hover {
|
| border-color: #007bff;
|
| background-color: #f0f8ff;
|
| }
|
| .btn {
|
| background-color: #007bff;
|
| color: white;
|
| padding: 10px 20px;
|
| border: none;
|
| border-radius: 5px;
|
| cursor: pointer;
|
| font-size: 16px;
|
| }
|
| .btn:hover {
|
| background-color: #0056b3;
|
| }
|
| .result {
|
| margin-top: 20px;
|
| padding: 20px;
|
| background-color: #f8f9fa;
|
| border-radius: 5px;
|
| border: 1px solid #dee2e6;
|
| }
|
| .json-output {
|
| background-color: #f4f4f4;
|
| padding: 15px;
|
| border-radius: 5px;
|
| font-family: monospace;
|
| white-space: pre-wrap;
|
| overflow-x: auto;
|
| max-height: 400px;
|
| overflow-y: auto;
|
| }
|
| .text-input {
|
| width: 100%;
|
| height: 100px;
|
| padding: 10px;
|
| border: 1px solid #ccc;
|
| border-radius: 5px;
|
| font-family: monospace;
|
| resize: vertical;
|
| }
|
| .tab-container {
|
| margin: 20px 0;
|
| }
|
| .tabs {
|
| display: flex;
|
| border-bottom: 1px solid #ccc;
|
| }
|
| .tab {
|
| padding: 10px 20px;
|
| cursor: pointer;
|
| border-bottom: 2px solid transparent;
|
| background-color: #f8f9fa;
|
| margin-right: 5px;
|
| }
|
| .tab.active {
|
| border-bottom-color: #007bff;
|
| background-color: white;
|
| }
|
| .tab-content {
|
| display: none;
|
| padding: 20px 0;
|
| }
|
| .tab-content.active {
|
| display: block;
|
| }
|
| </style>
|
| </head>
|
| <body>
|
| <div class="container">
|
| <div class="header">
|
| <h1>Document Text Extraction</h1>
|
| <p>Extract structured information from documents using AI</p>
|
| </div>
|
|
|
| <div class="tab-container">
|
| <div class="tabs">
|
| <div class="tab active" onclick="showTab('file')">Upload File</div>
|
| <div class="tab" onclick="showTab('text')">Enter Text</div>
|
| </div>
|
|
|
| <div id="file-tab" class="tab-content active">
|
| <form id="uploadForm" enctype="multipart/form-data">
|
| <div class="upload-area">
|
| <p>Choose a document to extract information</p>
|
| <p><small>Supported: PDF, DOCX, Images (PNG, JPG, etc.)</small></p>
|
| <input type="file" id="fileInput" name="file" accept=".pdf,.docx,.png,.jpg,.jpeg,.tiff,.bmp" style="margin: 10px 0;">
|
| <br>
|
| <button type="submit" class="btn">Extract Information</button>
|
| </div>
|
| </form>
|
| </div>
|
|
|
| <div id="text-tab" class="tab-content">
|
| <form id="textForm">
|
| <p>Enter text directly for information extraction:</p>
|
| <textarea id="textInput" class="text-input" placeholder="Enter document text here, e.g.: Invoice sent to John Doe on 01/15/2025 Invoice No: INV-1001 Amount: $1,500.00"></textarea>
|
| <br><br>
|
| <button type="submit" class="btn">Extract from Text</button>
|
| </form>
|
| </div>
|
| </div>
|
|
|
| <div id="result" class="result" style="display: none;">
|
| <h3>Extraction Results</h3>
|
| <div id="resultContent"></div>
|
| </div>
|
| </div>
|
|
|
| <script>
|
| function showTab(tabName) {
|
| // Hide all tab contents
|
| document.querySelectorAll('.tab-content').forEach(content => {
|
| content.classList.remove('active');
|
| });
|
|
|
| // Remove active class from all tabs
|
| document.querySelectorAll('.tab').forEach(tab => {
|
| tab.classList.remove('active');
|
| });
|
|
|
| // Show selected tab content
|
| document.getElementById(tabName + '-tab').classList.add('active');
|
|
|
| // Add active class to selected tab
|
| event.target.classList.add('active');
|
| }
|
|
|
| // File upload form handler
|
| document.getElementById('uploadForm').addEventListener('submit', async function(e) {
|
| e.preventDefault();
|
|
|
| const fileInput = document.getElementById('fileInput');
|
| if (!fileInput.files[0]) {
|
| alert('Please select a file');
|
| return;
|
| }
|
|
|
| const formData = new FormData();
|
| formData.append('file', fileInput.files[0]);
|
|
|
| try {
|
| showResult('Processing document, please wait...');
|
|
|
| const response = await fetch('/extract-from-file', {
|
| method: 'POST',
|
| body: formData
|
| });
|
|
|
| const result = await response.json();
|
| displayResult(result);
|
|
|
| } catch (error) {
|
| showResult('Error: ' + error.message);
|
| }
|
| });
|
|
|
| // Text form handler
|
| document.getElementById('textForm').addEventListener('submit', async function(e) {
|
| e.preventDefault();
|
|
|
| const text = document.getElementById('textInput').value;
|
| if (!text.trim()) {
|
| alert('Please enter some text');
|
| return;
|
| }
|
|
|
| try {
|
| showResult('Processing text, please wait...');
|
|
|
| const response = await fetch('/extract-from-text', {
|
| method: 'POST',
|
| headers: {
|
| 'Content-Type': 'application/json',
|
| },
|
| body: JSON.stringify({ text: text })
|
| });
|
|
|
| const result = await response.json();
|
| displayResult(result);
|
|
|
| } catch (error) {
|
| showResult('Error: ' + error.message);
|
| }
|
| });
|
|
|
| function showResult(message) {
|
| const resultDiv = document.getElementById('result');
|
| const contentDiv = document.getElementById('resultContent');
|
| contentDiv.innerHTML = message;
|
| resultDiv.style.display = 'block';
|
| }
|
|
|
| function displayResult(result) {
|
| let html = '';
|
|
|
| if (result.error) {
|
| html = `<div style="color: red;">Error: ${result.error}</div>`;
|
| } else {
|
| // Show structured data
|
| if (result.structured_data && Object.keys(result.structured_data).length > 0) {
|
| html += '<h4>Extracted Information:</h4>';
|
| html += '<table style="width: 100%; border-collapse: collapse; margin: 10px 0;">';
|
| html += '<tr style="background-color: #f8f9fa;"><th style="padding: 8px; border: 1px solid #dee2e6; text-align: left;">Field</th><th style="padding: 8px; border: 1px solid #dee2e6; text-align: left;">Value</th></tr>';
|
|
|
| for (const [key, value] of Object.entries(result.structured_data)) {
|
| html += `<tr><td style="padding: 8px; border: 1px solid #dee2e6; font-weight: bold;">${key}</td><td style="padding: 8px; border: 1px solid #dee2e6;">${value}</td></tr>`;
|
| }
|
| html += '</table>';
|
| } else {
|
| html += '<div style="color: orange;">No structured information found in the document.</div>';
|
| }
|
|
|
| // Show entities
|
| if (result.entities && result.entities.length > 0) {
|
| html += '<h4>Detected Entities:</h4>';
|
| html += '<div style="margin: 10px 0;">';
|
| result.entities.forEach(entity => {
|
| const confidence = Math.round(entity.confidence * 100);
|
| html += `<span style="display: inline-block; margin: 2px 4px; padding: 4px 8px; background-color: #e3f2fd; border: 1px solid #2196f3; border-radius: 15px; font-size: 12px;">
|
| ${entity.entity}: "${entity.text}" (${confidence}%)</span>`;
|
| });
|
| html += '</div>';
|
| }
|
|
|
| // Show raw JSON
|
| html += '<h4>Full Response:</h4>';
|
| html += `<div class="json-output">${JSON.stringify(result, null, 2)}</div>`;
|
| }
|
|
|
| showResult(html);
|
| }
|
| </script>
|
| </body>
|
| </html>
|
| """
|
| return html_content
|
|
|
|
|
| @app.get("/health")
|
| async def health_check():
|
| """Health check endpoint."""
|
| try:
|
| get_inference_pipeline()
|
| return {"status": "healthy", "message": "Model loaded successfully"}
|
| except Exception as e:
|
| return {"status": "unhealthy", "message": str(e)}
|
|
|
|
|
| @app.post("/extract-from-file")
|
| async def extract_from_file(file: UploadFile = File(...)):
|
| """Extract structured information from an uploaded file."""
|
| if not file:
|
| raise HTTPException(status_code=400, detail="No file provided")
|
|
|
|
|
| allowed_extensions = {'.pdf', '.docx', '.png', '.jpg', '.jpeg', '.tiff', '.bmp'}
|
| file_extension = Path(file.filename).suffix.lower()
|
|
|
| if file_extension not in allowed_extensions:
|
| raise HTTPException(
|
| status_code=400,
|
| detail=f"Unsupported file type: {file_extension}. Allowed: {', '.join(allowed_extensions)}"
|
| )
|
|
|
|
|
| with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as temp_file:
|
| shutil.copyfileobj(file.file, temp_file)
|
| temp_file_path = temp_file.name
|
|
|
| try:
|
|
|
| inference = get_inference_pipeline()
|
| result = inference.process_document(temp_file_path)
|
|
|
| return JSONResponse(content=result)
|
|
|
| except Exception as e:
|
| raise HTTPException(status_code=500, detail=str(e))
|
|
|
| finally:
|
|
|
| try:
|
| os.unlink(temp_file_path)
|
| except:
|
| pass
|
|
|
|
|
| @app.post("/extract-from-text")
|
| async def extract_from_text(request: Dict[str, str]):
|
| """Extract structured information from text."""
|
| text = request.get("text", "").strip()
|
|
|
| if not text:
|
| raise HTTPException(status_code=400, detail="No text provided")
|
|
|
| try:
|
|
|
| inference = get_inference_pipeline()
|
| result = inference.process_text_directly(text)
|
|
|
| return JSONResponse(content=result)
|
|
|
| except Exception as e:
|
| raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
| @app.get("/supported-formats")
|
| async def get_supported_formats():
|
| """Get list of supported file formats."""
|
| return {
|
| "supported_formats": [
|
| {"extension": ".pdf", "description": "PDF documents"},
|
| {"extension": ".docx", "description": "Microsoft Word documents"},
|
| {"extension": ".png", "description": "PNG images"},
|
| {"extension": ".jpg", "description": "JPEG images"},
|
| {"extension": ".jpeg", "description": "JPEG images"},
|
| {"extension": ".tiff", "description": "TIFF images"},
|
| {"extension": ".bmp", "description": "BMP images"}
|
| ],
|
| "entity_types": [
|
| "Name", "Date", "InvoiceNo", "Amount", "Address", "Phone", "Email"
|
| ]
|
| }
|
|
|
|
|
| @app.get("/model-info")
|
| async def get_model_info():
|
| """Get information about the loaded model."""
|
| try:
|
| inference = get_inference_pipeline()
|
| return {
|
| "model_path": inference.model_path,
|
| "model_name": inference.config.model_name,
|
| "max_length": inference.config.max_length,
|
| "entity_labels": inference.config.entity_labels,
|
| "num_labels": inference.config.num_labels
|
| }
|
| except Exception as e:
|
| raise HTTPException(status_code=503, detail=f"Model not loaded: {str(e)}")
|
|
|
|
|
| def main():
|
| """Run the FastAPI server."""
|
| print("Starting Document Text Extraction API Server...")
|
| print("Server will be available at: http://localhost:8000")
|
| print("Web interface: http://localhost:8000")
|
| print("API docs: http://localhost:8000/docs")
|
| print("Health check: http://localhost:8000/health")
|
|
|
| uvicorn.run(
|
| "api.app:app",
|
| host="0.0.0.0",
|
| port=8000,
|
| reload=True,
|
| log_level="info"
|
| )
|
|
|
|
|
| if __name__ == "__main__":
|
| main() |
