scripts/extract_huggingface_datasets.py

import os
import sys
import json
import re
import requests
from pathlib import Path

sys.stdout.reconfigure(encoding='utf-8')

# Output directories
BASE_OUT_DIR = Path(r"c:\Risu Solutions\ByteAstra\backend\data\ayurveda\scraped")
ASHTANGA_DIR = BASE_OUT_DIR / "ashtanga"
SUSHRUTA_DIR = BASE_OUT_DIR / "sushruta"
CHARAKA_DIR = BASE_OUT_DIR / "charaka"
RASAJALANIDHI_DIR = BASE_OUT_DIR / "rasajalanidhi"
JOURNALS_DIR = BASE_OUT_DIR / "journals"

# Ensure all output directories exist
for d in [ASHTANGA_DIR, SUSHRUTA_DIR, CHARAKA_DIR, RASAJALANIDHI_DIR, JOURNALS_DIR]:
    d.mkdir(parents=True, exist_ok=True)

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"
}

def slugify(text: str) -> str:
    text = text.lower().strip()
    text = re.sub(r"[^\w\s-]", "", text)
    text = re.sub(r"[\s_]+", "_", text)
    return text[:60]

def extract_ashtanga():
    print("\n=== Extracting Ashtanga Hridayam ===")
    url = "https://huggingface.co/datasets/vaishnavi0901/unsloth-adpt-ashtang_hridyam2-dataset/raw/main/unsloth-ashtang_hridyam_dataset2.json"
    
    try:
        resp = requests.get(url, headers=HEADERS, timeout=30)
        if resp.status_code != 200:
            print(f"Failed to fetch Ashtanga dataset: status {resp.status_code}")
            return
        
        data = resp.json()
        convs = data.get("conversations", [])
        print(f"Loaded {len(convs)} shlokas/conversations.")
        
        # Group by Sthanam and Chapter
        grouped = {}
        for idx, item in enumerate(convs):
            # item is a list: [human_msg, assistant_msg]
            if len(item) < 2:
                continue
            human_val = item[0].get("value", "").strip()
            assistant_val = item[1].get("value", "").strip()
            
            # Parse assistant JSON metadata
            try:
                meta = json.loads(assistant_val)
            except Exception:
                meta = {}
                
            # Parse metadata from the "source" field which contains comma-separated values
            source_val = meta.get("source", "")
            meta_parts = {}
            for part in source_val.split(","):
                if ":" in part:
                    k, v = part.split(":", 1)
                    meta_parts[k.strip().lower()] = v.strip()
                
            sthanam = meta_parts.get("sthanam", meta.get("Sthanam", meta.get("sthanam", "Sūtra-sthāna"))).strip()
            chapter = meta_parts.get("chapter", meta.get("Chapter", meta.get("chapter", "General"))).strip()
            label = meta_parts.get("label", meta.get("Label", meta.get("label", "General Principles"))).strip()
            # Clean trailing colons from label
            if label.endswith(":"):
                label = label[:-1].strip()
            explanation = meta.get("explanation", meta.get("Explanation", assistant_val)).strip()
            shloka_num = meta_parts.get("shloka number", meta.get("Shloka Number", meta.get("shloka_number", "")))
            
            key = (sthanam, chapter)
            grouped.setdefault(key, []).append({
                "shloka": human_val,
                "label": label,
                "explanation": explanation,
                "shloka_num": shloka_num
            })
            
        print(f"Grouped into {len(grouped)} chapters.")
        
        saved_count = 0
        for (sthanam, chapter), shlokas in grouped.items():
            sthanam_slug = slugify(sthanam)
            chapter_slug = slugify(str(chapter))
            filename = f"ashtanga_{sthanam_slug}_chapter_{chapter_slug}.md"
            filepath = ASHTANGA_DIR / filename
            
            # Create content
            lines = [
                "---",
                "source: Ashtanga Hridayam",
                f"chapter: Chapter {chapter}",
                f"section: {sthanam}",
                "---",
                "",
                f"# {sthanam} — Chapter {chapter}",
                ""
            ]
            
            for s_idx, s in enumerate(shlokas):
                lines.extend([
                    f"## {s['label']}",
                    "**Shloka:**",
                    "```",
                    s['shloka'],
                    "```",
                    "",
                    "**Explanation:**",
                    s['explanation'],
                    ""
                ])
                
            filepath.write_text("\n".join(lines), encoding="utf-8")
            saved_count += 1
            
        print(f"✓ Saved {saved_count} Ashtanga Hridayam chapter files.")
    except Exception as e:
        print("Error extracting Ashtanga:", e)

def extract_vedas_and_samhitas():
    print("\n=== Extracting Vedas, Charaka, Sushruta, Rasa Jala Nidhi, and Journals ===")
    url = "https://huggingface.co/datasets/shinigamiRaj/IndianVedasOriginal/resolve/main/continueousPreTrainData.jsonl"
    
    try:
        resp = requests.get(url, headers=HEADERS, stream=True, timeout=60)
        if resp.status_code != 200:
            print(f"Failed to fetch Vedas dataset: status {resp.status_code}")
            return
            
        line_idx = 0
        counts = {
            "charaka": 0,
            "sushruta": 0,
            "rasajalanidhi": 0,
            "journals": 0
        }
        
        for line in resp.iter_lines():
            if not line:
                continue
            line_idx += 1
            try:
                item = json.loads(line.decode("utf-8"))
                text = item.get("text", "").strip()
                if not text:
                    continue
                
                # Check collection tags
                if "[[ collection: charaka samhita" in text.lower():
                    # Parse Charaka
                    save_document(text, "Charaka Samhita", CHARAKA_DIR)
                    counts["charaka"] += 1
                elif "[[ collection: sushruta samhita" in text.lower():
                    # Parse Sushruta
                    save_document(text, "Sushruta Samhita", SUSHRUTA_DIR)
                    counts["sushruta"] += 1
                elif "[[ collection: rasa jala nidhi" in text.lower():
                    # Parse Rasa Jala Nidhi
                    save_document(text, "Rasa Jala Nidhi", RASAJALANIDHI_DIR)
                    counts["rasajalanidhi"] += 1
                elif "[[ collection: international research journal of ayurveda and yoga" in text.lower():
                    # Parse Journals
                    save_document(text, "IRJAY Journal", JOURNALS_DIR)
                    counts["journals"] += 1
            except Exception as parse_err:
                pass
                
        print(f"✓ Scan complete. Total lines read: {line_idx}")
        print(f"  - Extracted Charaka chapters/parts: {counts['charaka']}")
        print(f"  - Extracted Sushruta chapters/parts: {counts['sushruta']}")
        print(f"  - Extracted Rasa Jala Nidhi parts: {counts['rasajalanidhi']}")
        print(f"  - Extracted Journal articles: {counts['journals']}")
        
    except Exception as e:
        print("Error extracting Vedas:", e)

def save_document(text: str, source_name: str, out_dir: Path):
    # Parse headers from the first few lines of the text block
    lines = text.splitlines()
    title = ""
    section = ""
    chapter = ""
    
    # Simple heuristic to extract title, chapter and section
    for line in lines[:10]:
        line_clean = line.strip().strip("[]").replace("Collection: ", "").replace("Translator: ", "")
        if not line_clean:
            continue
        if "volume" in line_clean.lower() or "sthanam" in line_clean.lower() or "sthana" in line_clean.lower() or "vol." in line_clean.lower():
            section = line_clean
        elif "chapter " in line_clean.lower() or "chapter" in line_clean.lower() or "hymn" in line_clean.lower():
            chapter = line_clean
        elif len(line_clean) > 8 and not title and not line_clean.startswith("[[") and not "samhita" in line_clean.lower():
            title = line_clean

    if not chapter:
        chapter = "General"
    if not section:
        section = "General Section"
    if not title:
        title = chapter
        
    # Generate filename
    sec_slug = slugify(section)
    ch_slug = slugify(chapter)
    title_slug = slugify(title)
    
    filename = f"{sec_slug}_{ch_slug}_{title_slug}.md"
    # Ensure filename is unique if we write to same directory
    filepath = out_dir / filename
    
    # Strip the collection tag from top lines
    clean_lines = []
    for line in lines:
        if line.strip().startswith("[[") and "collection:" in line.lower():
            continue
        clean_lines.append(line)
        
    body = "\n".join(clean_lines).strip()
    
    if filepath.exists():
        existing_content = filepath.read_text(encoding="utf-8", errors="ignore")
        filepath.write_text(existing_content + "\n\n" + body, encoding="utf-8")
    else:
        frontmatter = (
            f"---\n"
            f"source: {source_name}\n"
            f"chapter: {chapter}\n"
            f"section: {section}\n"
            f"---\n\n"
        )
        filepath.write_text(frontmatter + body, encoding="utf-8")

if __name__ == "__main__":
    extract_ashtanga()
    extract_vedas_and_samhitas()