Update app.py

app.py

@@ -2,9 +2,11 @@ import gradio as gr
 import os
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import PyPDF2
-from io import BytesIO
 import torch
+from io import BytesIO
+from PIL import Image
+import fitz  # PyMuPDF
+from transformers import NougatProcessor, VisionEncoderDecoderModel
 
 # Set environment variables
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -65,30 +67,70 @@ AVAILABLE_MODELS = {
 current_model = None
 current_tokenizer = None
 current_model_name = None
+nougat_model = None
+nougat_processor = None
+
+@spaces.GPU(duration=200)
+def load_nougat_model():
+    """Load Nougat model for PDF processing"""
+    global nougat_model, nougat_processor
+    
+    if nougat_model is None or nougat_processor is None:
+        nougat_processor = NougatProcessor.from_pretrained("facebook/nougat-base")
+        nougat_model = VisionEncoderDecoderModel.from_pretrained("facebook/nougat-base")
+        nougat_model.to("cuda" if torch.cuda.is_available() else "cpu")
+    
+    return nougat_processor, nougat_model
 
+@spaces.GPU(duration=200)
 def extract_text_from_pdf(pdf_bytes):
-    """Extract text from uploaded PDF file"""
+    """Extract text from uploaded PDF file using Nougat"""
     if pdf_bytes is None:
         return default_paper_content
     
     try:
-        # Ensure pdf_bytes is bytes type
-        if isinstance(pdf_bytes, str):
-            return pdf_bytes  # If already a string, return directly
+        # Load Nougat model
+        processor, model = load_nougat_model()
+        
+        # Convert PDF to images
+        pdf_document = fitz.open(stream=pdf_bytes, filetype="pdf")
+        full_text = ""
+        
+        for page_num in range(len(pdf_document)):
+            page = pdf_document.load_page(page_num)
+            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x zoom for better quality
+            
+            # Convert to PIL Image
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            
+            # Process with Nougat
+            pixel_values = processor(img, return_tensors="pt").pixel_values.to(model.device)
+            
+            # Generate text
+            outputs = model.generate(
+                pixel_values,
+                min_length=1,
+                max_new_tokens=1024,  # Adjust based on expected page content length
+                bad_words_ids=[[processor.tokenizer.unk_token_id]],
+            )
+            
+            # Decode and post-process
+            page_text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+            page_text = processor.post_process_generation(page_text, fix_markdown=True)
             
-        # Use bytes object directly
-        pdf_reader = PyPDF2.PdfReader(BytesIO(pdf_bytes))
+            full_text += page_text + "\n\n"
         
-        # Extract text from all pages
-        text = ""
-        for page_num in range(len(pdf_reader.pages)):
-            page = pdf_reader.pages[page_num]
-            text += page.extract_text() + "\n\n"
+        # Clear GPU memory
+        del pixel_values, outputs
+        torch.cuda.empty_cache()
         
-        return text
+        return full_text
     except Exception as e:
         print(f"PDF extraction error: {str(e)}")
         return default_paper_content
+    finally:
+        # Clear GPU memory
+        torch.cuda.empty_cache()
 
 def load_model(model_name):
     """Load model and tokenizer on demand"""