Spaces:
Running
Running
File size: 2,870 Bytes
c8c7a9e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import os
from typing import List, Dict, Any
import fitz # PyMuPDF
import docx
import requests
import io
def process_text_file(file_url: str) -> List[Dict[str, Any]]:
_, extension = os.path.splitext(file_url)
extension = extension.lower()
if "?alt=media&token=" in extension:
extension = list(extension.split("?"))[0]
if extension == '.txt':
return process_txt(file_url)
elif extension == '.pdf':
return process_pdf(file_url)
elif extension == '.docx':
return process_docx(file_url)
else:
raise ValueError(f"Unsupported text file type: {extension}")
def process_txt(txt_url: str) -> List[Dict[str, Any]]:
# Fetch the TXT file content from the URL
response = requests.get(txt_url)
# Check if the request was successful
if response.status_code == 200:
content = response.text
return [{
"file_name": os.path.basename(txt_url),
"text": content,
"page_number": 1
}]
else:
print(f"Failed to fetch the TXT file. Status code: {response.status_code}")
return []
def process_pdf(pdf_url: str) -> List[Dict[str, Any]]:
# Fetch the PDF file content from the URL
response = requests.get(pdf_url)
# Check if the request was successful
if response.status_code == 200:
# Load the PDF file from the response content
pdf_stream = io.BytesIO(response.content)
# Open the PDF file with PyMuPDF
pdf_document = fitz.open(stream=pdf_stream, filetype="pdf")
# Extract text from all pages
pdf_text = ""
for page_num in range(len(pdf_document)):
page = pdf_document.load_page(page_num) # Load the page
pdf_text += page.get_text("text") # Extract text from the page
return [{
"file_name": os.path.basename(pdf_url),
"text": pdf_text
}]
else:
print(f"Failed to fetch the PDF file. Status code: {response.status_code}")
return []
def process_docx(docx_url: str) -> List[Dict[str, Any]]:
# Fetch the DOCX file content from the URL
response = requests.get(docx_url)
# Check if the request was successful
if response.status_code == 200:
# Load the DOCX file from the response content
docx_stream = io.BytesIO(response.content)
# Open the DOCX file with python-docx
doc = docx.Document(docx_stream)
# Extract text from the DOCX file
content = "\n".join([para.text for para in doc.paragraphs])
return [{
"file_name": os.path.basename(docx_url),
"text": content,
"page_number": 1 # DOCX doesn't have pages, so just 1
}]
else:
print(f"Failed to fetch the DOCX file. Status code: {response.status_code}")
|