test build

backend/scripts/parse_wikipedia.py

@@ -0,0 +1,110 @@
+import bz2
+import xml.etree.ElementTree as ET
+import json
+import os
+from pathlib import Path
+from typing import Generator, Dict, Any
+
+def extract_text_from_wikipedia_dump(bz2_path: str, output_dir: str, limit: int = None) -> int:
+    """Extract text from Wikipedia XML dump"""
+    os.makedirs(output_dir, exist_ok=True)
+    
+    count = 0
+    with bz2.BZ2File(bz2_path, 'r') as f:
+        context = ET.iterparse(f, events=('start', 'end'))
+        context = iter(context)
+        
+        event, root = next(context)
+        
+        for event, elem in context:
+            if event == 'end' and elem.tag.endswith('page'):
+                doc = parse_page(elem)
+                if doc and doc.get('text'):
+                    filename = f"{output_dir}/wiki_{doc['id']}.json"
+                    with open(filename, 'w', encoding='utf-8') as out:
+                        json.dump(doc, out, ensure_ascii=False)
+                    
+                    count += 1
+                    if count % 1000 == 0:
+                        print(f"Processed {count} documents")
+                    
+                    if limit and count >= limit:
+                        break
+                
+                root.clear()
+    
+    return count
+
+def parse_page(page_elem) -> Dict[str, Any]:
+    """Parse a single Wikipedia page"""
+    ns = {'wiki': 'http://www.mediawiki.org/xml/export-0.10/'}
+    
+    title_elem = page_elem.find('wiki:title', ns) or page_elem.find('.//title')
+    page_id_elem = page_elem.find('wiki:id', ns) or page_elem.find('.//id')
+    revision_elem = page_elem.find('.//revision')
+    
+    if revision_elem is None:
+        return None
+    
+    text_elem = revision_elem.find('wiki:text', ns) or revision_elem.find('.//text')
+    
+    if title_elem is None or text_elem is None:
+        return None
+    
+    title = title_elem.text or ""
+    page_id = page_id_elem.text if page_id_elem is not None else ""
+    text = text_elem.text or ""
+    
+    # Skip redirects and special pages
+    if text.startswith('#REDIRECT') or title.startswith('Wikipedia:'):
+        return None
+    
+    # Basic text cleaning
+    text = clean_wikipedia_markup(text)
+    
+    if len(text.split()) < 50:
+        return None
+    
+    return {
+        "id": f"wiki_{page_id}",
+        "title": title,
+        "text": text,
+        "source": "wikipedia",
+        "lang": "ru"  # Change for English
+    }
+
+def clean_wikipedia_markup(text: str) -> str:
+    """Remove Wikipedia markup"""
+    import re
+    
+    # Remove wiki links
+    text = re.sub(r'\[\[([^\|]*\|)?([^\]]*)\]\]', r'\2', text)
+    # Remove external links
+    text = re.sub(r'\[http[^\s]+ ([^\]]*)\]', r'\1', text)
+    # Remove templates
+    text = re.sub(r'\{\{[^}]*\}\}', '', text)
+    # Remove categories
+    text = re.sub(r'\[\[Category:[^\]]*\]\]', '', text)
+    # Remove refs
+    text = re.sub(r'<ref[^>]*>[^<]*</ref>', '', text)
+    # Remove HTML tags
+    text = re.sub(r'<[^>]*>', '', text)
+    # Clean whitespace
+    text = re.sub(r'\n\n+', '\n', text)
+    text = re.sub(r' +', ' ', text)
+    
+    return text.strip()
+
+if __name__ == "__main__":
+    import sys
+    
+    if len(sys.argv) < 3:
+        print("Usage: python parse_wikipedia.py <input_bz2> <output_dir> [limit]")
+        sys.exit(1)
+    
+    input_file = sys.argv[1]
+    output_dir = sys.argv[2]
+    limit = int(sys.argv[3]) if len(sys.argv) > 3 else None
+    
+    count = extract_text_from_wikipedia_dump(input_file, output_dir, limit)
+    print(f"Successfully processed {count} documents")