Update

app.py

@@ -2,16 +2,12 @@ from pdfminer.high_level import extract_pages
 from pdfminer.layout import LTTextContainer
 from tqdm import tqdm
 import re
-from zipfile import ZipFile
+import io
+import zipfile
 import gradio as gr
 import os
 from llama_cpp import Llama
-from gpt4all import GPT4All
-import transformers
-# from transformers import GemmaTokenizer, AutoModelForCausalLM
-# from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-import accelerate
-import torch
+import tempfile
 
 def process_document(pdf_path, page_ids=None):
 
@@ -112,33 +108,32 @@ def deidentify_doc(pdftext="", prompt="", maxtokens=600, temperature=0.0001, top
         top_p=top_probability,
     )
     output = outputs[0]["generated_text"][len(prompt):]
-    
+
     return output
 
+def mkdir(dir):
+    if not os.path.exists(dir):
+        os.makedirs(dir)
+
 def pdf_to_text(files, prompt="", maxtokens=600, temperature=1.2, top_probability=0.95):
-    output_files = []
-    for file in files:
-        file_name = os.path.basename(file)
-        file_name_splt = file_name.split('.')
-        print('File name is ', file_name)
-        if(len(file_name_splt)>1 and file_name_splt[1]=='pdf'):
-           page2content = process_document(file, page_ids=[0])
-           pdftext = page2content[1]
-           if(pdftext):
-               anonymized_text = deidentify_doc(pdftext, prompt, maxtokens, temperature, top_probability)
-               # html = txt_to_html(display_text)
-               # with open('out.html', "w", encoding="utf-8") as file:
-               #     file.write(html)
-               with open(file_name_splt[0]+'.txt', 'w') as outputfile:
-                # Write some text to the file
-                outputfile.write(anonymized_text)
-                output_files.append(file_name_splt[0]+'.txt')
-    zipf = ZipFile('anonymized_reports', 'w')
-    for file in output_files:
-        zipf.write(file, os.path.basename(file))
-    return 'anonymized_reports'
-
-# model_id = "Meta-Llama-3-8B-Instruct.Q5_K_M.gguf"
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for file in files:
+            file_name = os.path.basename(file)
+            file_name_splt = file_name.split('.')
+            if(len(file_name_splt)>1 and file_name_splt[1]=='pdf'):
+               page2content = process_document(file, page_ids=[0])
+               pdftext = page2content[1]
+               if(pdftext):
+                   anonymized_text = deidentify_doc(pdftext, prompt, maxtokens, temperature, top_probability)
+                   zf.writestr(file_name_splt[0]+'.txt', anonymized_text)
+    zip_buffer.seek(0)
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.zip') as temp_file:
+        temp_file.write(zip_buffer.getvalue())
+        temp_file_path = temp_file.name
+    return temp_file_path
+
+# model_id = "D:/llama/meta-llama/Meta-Llama-3-8B-Instruct.Q5_K_M.gguf"
 # model = Llama(model_path=model_id, n_ctx=2048, n_threads=8, n_gpu_layers=32, n_batch=64)
 
 # model = GPT4All("Meta-Llama-3-8B-Instruct.Q4_0.gguf", n_threads=8, device='gpu')