Upload 2 files

app.py

@@ -2,9 +2,6 @@ import os
 import gradio as gr
 import main
 
-#os.environ["CUDA_VISIBLE_DEVICES"]='0'
-#os.environ["USE_GPU"]="True"
-
 
 def predict_from_pdf(pdf_file):
     upload_dir = "./catalogue/"
@@ -35,9 +32,9 @@ demo = gr.Interface(
     outputs=["json", "text"],
     examples=pdf_examples,
     title="Open Source PDF Catalog Parser",
-    description="Efficient PDF catalog processing using PyMuPDF and OpenLLM",
+    description="Efficient PDF catalog processing using MinerU and OpenLLM",
     article="Uses MinerU for layout analysis and DeepSeek-7B for structured extraction"
 )
 
 if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)

main.py

@@ -42,23 +42,26 @@ class PDFProcessor:
         self.output_dir.mkdir(exist_ok=True)
 
     def _initialize_emb_model(self, model_name):
-        # model = SentenceTransformer("sentence-transformers/" + model_name)
-        # model = SentenceTransformer(model_name)
-        # model.save('models/'+ model_name)
-        # Load model directly
-        from transformers import AutoTokenizer, AutoModel
+        try:
+            model = SentenceTransformer("sentence-transformers/" + model_name)
+            model.save('models/'+ model_name)
+            return model
+        except:
+            # Load model directly
+            from transformers import AutoTokenizer, AutoModel
 
-        tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-        model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-        return model
+            tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+            model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+            return model
 
     def _initialize_llm(self, model_name):
         """Initialize LLM with automatic download if needed"""
+        """
         model_path = os.path.join("models/", model_name)
         if os.path.exists(model_path):
             return Llama(
                 model_path=model_path,
-                n_ctx=4096,
+                n_ctx=2048,
                 n_gpu_layers=35 if os.getenv('USE_GPU') else 0,
                 n_threads=os.cpu_count() - 1,
                 verbose=False
@@ -67,11 +70,16 @@ class PDFProcessor:
             return Llama.from_pretrained(
                 repo_id="TheBloke/deepseek-llm-7B-base-GGUF",
                 filename=model_name,
-                n_ctx=4096,
+                n_ctx=2048,
                 n_threads=os.cpu_count() - 1,
                 n_gpu_layers=35 if os.getenv('USE_GPU') else 0,
                 verbose=False
             )
+        """
+        # Load model directly
+        from transformers import AutoModel
+        model = AutoModel.from_pretrained("TheBloke/deepseek-llm-7B-base-GGUF")
+        return model
     
     def process_pdf(self, pdf_path: str) -> Dict:
         """Process PDF using MinerU pipeline"""