Update app.py

app.py

@@ -1,82 +1,177 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from PIL import Image
+from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
 import torch
+from qwen_vl_utils import process_vision_info
+import fitz  # PyMuPDF
+import io
+import os
+from typing import List, Dict
+import tempfile
+import uvicorn
 
-app = FastAPI()
-
-MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-print(f"Loading model on {device}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    device_map="auto",
-    trust_remote_code=True
+app = FastAPI(
+    title="PDF OCR API",
+    description="Convert PDF to images and extract text using Qari OCR model",
+    version="1.0.0"
 )
-print("Model loaded!")
 
-class ChatRequest(BaseModel):
-    message: str
-    max_tokens: int = 512
-    temperature: float = 0.7
+# Initialize model and processor
+model_name = "NAMAA-Space/Qari-OCR-0.2.2.1-VL-2B-Instruct"
+model = None
+processor = None
+max_tokens = 2000
 
-class CompletionRequest(BaseModel):
-    messages: list
-    max_tokens: int = 512
-    temperature: float = 0.7
-    stream: bool = False
+@app.on_event("startup")
+async def load_model():
+    """Load model on startup"""
+    global model, processor
+    model = Qwen2VLForConditionalGeneration.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto"
+    )
+    processor = AutoProcessor.from_pretrained(model_name)
+    print("Model loaded successfully!")
+
+def pdf_to_images(pdf_bytes: bytes) -> List[Image.Image]:
+    """Convert PDF pages to PIL Images"""
+    images = []
+    pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
+    
+    for page_num in range(len(pdf_document)):
+        page = pdf_document[page_num]
+        # Render page to image at 300 DPI for better quality
+        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        images.append(img)
+    
+    pdf_document.close()
+    return images
 
-@app.post("/chat")
-def chat(req: ChatRequest):
-    messages = [{"role": "user", "content": req.message}]
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer([text], return_tensors="pt").to(device)
+def process_image_ocr(image: Image.Image, temp_path: str) -> str:
+    """Process a single image with OCR"""
+    prompt = "Below is the image of one page of a document, as well as some raw textual content that was previously extracted for it. Just return the plain text representation of this document as if you were reading it naturally. Do not hallucinate."
+    
+    # Save image temporarily
+    image.save(temp_path)
+    
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": f"file://{temp_path}"},
+                {"type": "text", "text": prompt},
+            ],
+        }
+    ]
     
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=req.max_tokens,
-        temperature=req.temperature,
-        do_sample=True
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
     )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to(model.device)
+    
+    generated_ids = model.generate(**inputs, max_new_tokens=max_tokens)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("assistant\n")[-1].strip()
+    # Clean up temp file
+    if os.path.exists(temp_path):
+        os.remove(temp_path)
     
-    return {"response": response}
+    return output_text
 
-@app.post("/v1/chat/completions")
-def completions(req: CompletionRequest):
-    text = tokenizer.apply_chat_template(req.messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer([text], return_tensors="pt").to(device)
+@app.post("/ocr", response_model=Dict)
+async def extract_text_from_pdf(file: UploadFile = File(...)):
+    """
+    Extract text from PDF file using OCR
     
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=req.max_tokens,
-        temperature=req.temperature,
-        do_sample=True
-    )
+    - **file**: PDF file (max 5MB)
+    
+    Returns JSON with extracted text for each page
+    """
+    # Validate file type
+    if not file.filename.endswith('.pdf'):
+        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
+    
+    # Read file content
+    pdf_bytes = await file.read()
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("assistant\n")[-1].strip()
+    # Check file size (5MB limit)
+    file_size = len(pdf_bytes)
+    if file_size > 5 * 1024 * 1024:  # 5MB in bytes
+        raise HTTPException(
+            status_code=400, 
+            detail=f"File size ({file_size / (1024*1024):.2f}MB) exceeds 5MB limit"
+        )
     
+    try:
+        # Convert PDF to images
+        images = pdf_to_images(pdf_bytes)
+        
+        if not images:
+            raise HTTPException(status_code=400, detail="No pages found in PDF")
+        
+        # Process each page
+        results = []
+        with tempfile.TemporaryDirectory() as temp_dir:
+            for i, image in enumerate(images):
+                temp_path = os.path.join(temp_dir, f"page_{i}.png")
+                try:
+                    ocr_text = process_image_ocr(image, temp_path)
+                    results.append({
+                        "page": i + 1,
+                        "text": ocr_text
+                    })
+                except Exception as e:
+                    results.append({
+                        "page": i + 1,
+                        "error": str(e)
+                    })
+        
+        return {
+            "success": True,
+            "filename": file.filename,
+            "total_pages": len(images),
+            "file_size_mb": round(file_size / (1024*1024), 2),
+            "results": results
+        }
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")
+
+@app.get("/")
+async def root():
+    """API status endpoint"""
     return {
-        "choices": [{
-            "message": {"role": "assistant", "content": response},
-            "finish_reason": "stop"
-        }]
+        "message": "PDF OCR API is running",
+        "model": model_name,
+        "max_file_size": "5MB",
+        "endpoints": {
+            "POST /ocr": "Extract text from PDF"
+        }
     }
 
 @app.get("/health")
-def health():
-    return {"status": "ok"}
-
-
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None
+    }
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)
-
- 
+    uvicorn.run(app, host="0.0.0.0", port=7860)