Spaces:
Running
Running
File size: 10,002 Bytes
43cd37c |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 |
# Media_Wiki.py
# Description: This file contains the functions to import MediaWiki dumps into the media_db and Chroma databases.
#######################################################################################################################
#
# Imports
import json
import logging
import os
import re
import traceback
from typing import List, Dict, Any, Iterator, Optional
# 3rd-Party Imports
import mwparserfromhell
import mwxml
import yaml
#
# Local Imports
from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords
from App_Function_Libraries.RAG.ChromaDB_Library import process_and_store_content
#
#######################################################################################################################
#
# Functions:
# Load configuration
def load_mediawiki_import_config():
config_path = os.path.join(os.path.dirname(__file__), '..', '..', 'Config_Files', 'mediawiki_import_config.yaml')
with open(config_path, 'r') as f:
return yaml.safe_load(f)
config = load_mediawiki_import_config()
def setup_logger(name: str, level: int = logging.INFO, log_file: Optional[str] = None) -> logging.Logger:
"""Set up and return a logger with the given name and level."""
logger = logging.getLogger(name)
logger.setLevel(level)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
if log_file:
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
return logger
# Usage
logger = setup_logger('mediawiki_import', log_file='mediawiki_import.log')
# End of setup
#######################################################################################################################
#
# Functions:
def parse_mediawiki_dump(file_path: str, namespaces: List[int] = None, skip_redirects: bool = False) -> Iterator[
Dict[str, Any]]:
dump = mwxml.Dump.from_file(open(file_path, encoding='utf-8'))
for page in dump.pages:
if skip_redirects and page.redirect:
continue
if namespaces and page.namespace not in namespaces:
continue
for revision in page:
wikicode = mwparserfromhell.parse(revision.text)
plain_text = wikicode.strip_code()
yield {
"title": page.title,
"content": plain_text,
"namespace": page.namespace,
"page_id": page.id,
"revision_id": revision.id,
"timestamp": revision.timestamp
}
logger.debug(f"Yielded page: {page.title}")
def optimized_chunking(text: str, chunk_options: Dict[str, Any]) -> List[Dict[str, Any]]:
sections = re.split(r'\n==\s*(.*?)\s*==\n', text)
chunks = []
current_chunk = ""
current_size = 0
logging.debug(f"optimized_chunking: Processing text with {len(sections) // 2} sections")
for i in range(0, len(sections), 2):
section_title = sections[i] if i > 0 else "Introduction"
section_content = sections[i + 1] if i + 1 < len(sections) else ""
if current_size + len(section_content) > chunk_options['max_size']:
if current_chunk:
chunks.append({"text": current_chunk, "metadata": {"section": section_title}})
current_chunk = section_content
current_size = len(section_content)
else:
current_chunk += f"\n== {section_title} ==\n" + section_content
current_size += len(section_content)
if current_chunk:
chunks.append({"text": current_chunk, "metadata": {"section": "End"}})
return chunks
def process_single_item(content: str, title: str, wiki_name: str, chunk_options: Dict[str, Any],
is_combined: bool = False, item: Dict[str, Any] = None, api_name: str = None):
try:
logging.debug(f"process_single_item: Processing item: {title}")
# Create a unique URL using the wiki name and article title
encoded_title = title.replace(" ", "_")
url = f"mediawiki:{wiki_name}:{encoded_title}"
logging.debug(f"Generated URL: {url}")
result = add_media_with_keywords(
url=url, # Use the generated URL here
title=title,
media_type="mediawiki_dump" if is_combined else "mediawiki_article",
content=content,
keywords=f"mediawiki,{wiki_name}" + (",full_dump" if is_combined else ",article"),
prompt="",
summary="",
transcription_model="",
author="MediaWiki",
ingestion_date=item['timestamp'].strftime('%Y-%m-%d') if item else None
)
logging.debug(f"Result from add_media_with_keywords: {result}")
# Unpack the result
media_id, message = result
logging.info(f"Media item result: {message}")
logging.debug(f"Final media_id: {media_id}")
chunks = optimized_chunking(content, chunk_options)
for i, chunk in enumerate(chunks):
logging.debug(f"Processing chunk {i + 1}/{len(chunks)} for item: {title}")
# FIXME
# def process_and_store_content(content: str, collection_name: str, media_id: int, file_name: str,
# create_embeddings: bool = False, create_summary: bool = False,
# api_name: str = None):
if api_name:
process_and_store_content(chunk['text'], f"mediawiki_{wiki_name}", media_id, title, True, True, api_name)
else:
process_and_store_content(chunk['text'], f"mediawiki_{wiki_name}", media_id, title)
logging.info(f"Successfully processed item: {title}")
except Exception as e:
logging.error(f"Error processing item {title}: {str(e)}")
logging.error(f"Exception details: {traceback.format_exc()}")
def load_checkpoint(file_path: str) -> int:
if os.path.exists(file_path):
with open(file_path, 'r') as f:
return json.load(f)['last_processed_id']
return 0
def save_checkpoint(file_path: str, last_processed_id: int):
with open(file_path, 'w') as f:
json.dump({'last_processed_id': last_processed_id}, f)
def import_mediawiki_dump(
file_path: str,
wiki_name: str,
namespaces: List[int] = None,
skip_redirects: bool = False,
chunk_options: Dict[str, Any] = None,
single_item: bool = False,
progress_callback: Any = None,
api_name: str = None,
api_key: str = None
) -> Iterator[str]:
try:
logging.info(f"Importing MediaWiki dump: {file_path}")
if chunk_options is None:
chunk_options = config['chunking']
checkpoint_file = f"{wiki_name}_import_checkpoint.json"
last_processed_id = load_checkpoint(checkpoint_file)
total_pages = count_pages(file_path, namespaces, skip_redirects)
processed_pages = 0
yield f"Found {total_pages} pages to process."
for item in parse_mediawiki_dump(file_path, namespaces, skip_redirects):
if item['page_id'] <= last_processed_id:
continue
# FIXME - ensure this works...
if api_name is not None:
# FIXME - add API key to the call/params
process_single_item(item['content'], item['title'], wiki_name, chunk_options, False, item, api_name)
process_single_item(item['content'], item['title'], wiki_name, chunk_options, False, item)
save_checkpoint(checkpoint_file, item['page_id'])
processed_pages += 1
if progress_callback is not None:
progress_callback(processed_pages / total_pages, f"Processed page: {item['title']}")
yield f"Processed page {processed_pages}/{total_pages}: {item['title']}"
os.remove(checkpoint_file) # Remove checkpoint file after successful import
yield f"Successfully imported and indexed MediaWiki dump: {wiki_name}"
except FileNotFoundError:
logger.error(f"MediaWiki dump file not found: {file_path}")
yield f"Error: File not found - {file_path}"
except PermissionError:
logger.error(f"Permission denied when trying to read: {file_path}")
yield f"Error: Permission denied - {file_path}"
except Exception as e:
logger.exception(f"Error during MediaWiki import: {str(e)}")
yield f"Error during import: {str(e)}"
def count_pages(file_path: str, namespaces: List[int] = None, skip_redirects: bool = False) -> int:
"""
Count the number of pages in a MediaWiki XML dump file.
Args:
file_path (str): Path to the MediaWiki XML dump file.
namespaces (List[int], optional): List of namespace IDs to include. If None, include all namespaces.
skip_redirects (bool, optional): Whether to skip redirect pages.
Returns:
int: The number of pages in the dump file.
"""
try:
dump = mwxml.Dump.from_file(open(file_path, encoding='utf-8'))
count = 0
for page in dump.pages:
if skip_redirects and page.redirect:
continue
if namespaces and page.namespace not in namespaces:
continue
count += 1
return count
except Exception as e:
logger.error(f"Error counting pages in MediaWiki dump: {str(e)}")
return 0
#
# End of Media_Wiki.py
#######################################################################################################################
|