Spaces:
Paused
Paused
| import json | |
| from pathlib import Path | |
| from time import perf_counter | |
| from typing import Any, Dict | |
| from tqdm.auto import tqdm | |
| def folder_to_json(folder_in: Path, json_path: Path): | |
| """ | |
| Process JSON lines from files in a given folder and write processed data to a new JSON file. | |
| Parameters: | |
| folder_in (Path): Path to the input folder containing the JSON files to process. | |
| json_path (Path): Path to the output JSON file where the processed data will be written. | |
| Example: | |
| folder_to_json(Path("/path/to/input/folder"), Path("/path/to/output.json")) | |
| """ | |
| json_out = [] # Initialize list to hold processed JSON data from all files | |
| process_start = perf_counter() | |
| # Use rglob to get all JSON files and sort them by their full path | |
| all_files = sorted(folder_in.rglob('*wiki*'), key=lambda x: str(x)) | |
| # Initialize progress bar with total file count, description, and unit of progress | |
| with tqdm(total=len(all_files), desc='Processing', unit='file') as pbar: | |
| # Iterate through all files in the input folder in order | |
| for file_path in all_files: | |
| # Update progress bar postfix to display current file and directory | |
| pbar.set_postfix_str(f"File: {file_path.name} | Dir: {file_path.parent}", refresh=True) | |
| # Open and read the current file | |
| with open(file_path, 'r', encoding='utf-8') as f: | |
| for line in f: | |
| # Load JSON data from each line and process it | |
| article = json.loads(line) | |
| # Add restructured article to the output list | |
| json_out.extend([restructure_articles(article)]) | |
| # Update progress bar after processing each file | |
| pbar.update(1) | |
| time_taken_to_process = perf_counter() - process_start | |
| pbar.write(f"Wiki processed in {round(time_taken_to_process, 2)} seconds!") | |
| # Notify that the writing process is starting | |
| pbar.write("Writing file!") | |
| write_start = perf_counter() | |
| # Open the output file and write the processed data as JSON | |
| with open(json_path, "w", encoding='utf-8') as outfile: | |
| json.dump(json_out, outfile) | |
| time_taken_to_write = perf_counter() - write_start | |
| # Notify that the writing process is complete | |
| pbar.write(f"File written in {round(time_taken_to_write, 2)} seconds!") | |
| def restructure_articles(article: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Restructures the given article into haystack's format, separating content and meta data. | |
| Args: | |
| - article (Dict[str, Any]): The article to restructure. | |
| Returns: | |
| - Dict[str, Any]: The restructured article. | |
| """ | |
| # Extract content and separate meta data | |
| article_out = { | |
| 'content': article['text'], | |
| 'meta': {k: v for k, v in article.items() if k != 'text'} | |
| } | |
| return article_out | |
| if __name__ == '__main__': | |
| proj_dir = Path(__file__).parents[2] | |
| folder = proj_dir / 'data/raw/output' | |
| file_out = proj_dir / 'data/consolidated/simple_wiki.json' | |
| folder_to_json(folder, file_out) | |
| print('Done!') | |