multiple pdfs

app.py

@@ -68,28 +68,39 @@ def clean(text):
     return text
 
 
-def pdf2dataset(file, progress=gr.Progress()):
+def pdf2dataset(pathes, progress=gr.Progress()):
     progress(0, desc="Starting...")
-    reader = PdfReader(file)
-    num_pages = len(reader.pages)
-    dataset_name = f"{random.getrandbits(128):x}"
+    readers = []
+    for path in pathes:
+        if not path.endswith(".pdf"):
+            raise gr.Error(f"Failed to read {path.split('/')[-1]}.")
+    readers = [PdfReader(path) for path in pathes]
+    num_pages = sum(len(reader.pages) for reader in readers)
+    filenames = [path.split("/")[-1] for path in pathes]
+
+    # Convert the PDFs to text
     page_texts = []
-
-    for page in progress.tqdm(reader.pages, total=num_pages, desc="Converting pages"):
-        page_text = page.extract_text()
-        page_text = clean(page_text)
-        page_texts.append(page_text)
-
+    page_filenames = []
+    progress(0, desc="Converting pages...")
+    for reader, filename in zip(readers, filenames):
+        for page in reader.pages:
+            page_text = page.extract_text()
+            page_text = clean(page_text)
+            page_texts.append(page_text)
+            page_filenames.append(filename)
+            progress(len(page_texts) / num_pages, desc="Converting pages...")
+
+    # Upload the dataset to Hugging Face
     progress(0, desc="Uploading to Hugging Face...")
-    dataset = Dataset.from_dict({"text": page_texts})
-    dataset.push_to_hub(f"pdf2dataset/{dataset_name}", token=os.getenv("TOKEN"))
+    dataset = Dataset.from_dict({"text": page_texts, "source": page_filenames})
+    dataset_name = f"{random.getrandbits(128):x}"
+    dataset.push_to_hub(f"pdf2dataset/{dataset_name}", token=os.getenv("HF_TOKEN"))
     progress(1, desc="Done!")
 
-    instrctions = instructions_template.substitute(dataset_name=dataset_name)
-    preview = dataset["text"][:10]
-    preview = pd.DataFrame(preview, columns=["text"])
+    instructions = instructions_template.substitute(dataset_name=dataset_name)
+    preview = pd.DataFrame(dataset[:10])
     print(f"Dataset {dataset_name} uploaded successfully.")
-    return instrctions, preview, dataset_name
+    return instructions, preview, dataset_name
 
 
 def delete_dataset(dataset_name):
@@ -130,17 +141,17 @@ dataset = load_dataset("pdf2dataset/$dataset_name")
 
 with gr.Blocks() as demo:
     gr.Markdown("# PDF to 🤗 Dataset")
-    gr.Markdown("## 1️⃣ Upload a PDF")
-    file = gr.File(file_types=["pdf"], height=50)
+    gr.Markdown("## 1️⃣ Upload PDFs")
+    file = gr.File(file_types=["pdf"], file_count="multiple")
     gr.Markdown(caution_text)
-    gr.Markdown("## 2️⃣ Convert the PDF and upload")
+    gr.Markdown("## 2️⃣ Convert the PDFs and upload")
     convert_button = gr.Button("🔄 Convert and upload")
     preview = gr.Dataframe(
-        label="Preview (first 10 rows)", headers=["text"], datatype=["str"], row_count=10, wrap=True, height=200
+        label="Preview (first 10 rows)", headers=["text", "source"], datatype=["str", "str"], row_count=10, wrap=True, height=200
     )
     gr.Markdown("## 3️⃣ Use the dataset in your code")
     instructions = gr.Markdown(instructions_template.substitute(dataset_name="generated_dataset_name"))
-    gr.Markdown("## 4️⃣ Delete the (optional)")
+    gr.Markdown("## 4️⃣ Delete the dataset (optional)")
     dataset_name_to_delete = gr.Textbox("", placeholder="Enter dataset name to delete", label="Dataset to delete")
     delete_button = gr.Button("🗑️ Delete dataset")