datasets.py

from __future__ import annotations
import os
import pathlib
import requests
from typing import Iterable, List, Tuple


SUPPORTED = {
    "FB15k-237": {
        "bases": [
            "https://raw.githubusercontent.com/ibalazevic/Tucker/master/data/FB15k-237",
            "https://raw.githubusercontent.com/TimDettmers/ConvE/master/FB15k-237",
        ],
        "files": ["train.txt", "valid.txt", "test.txt"],
    },
    "WN18RR": {
        "bases": [
            "https://raw.githubusercontent.com/ibalazevic/Tucker/master/data/WN18RR",
            "https://raw.githubusercontent.com/TimDettmers/ConvE/master/WN18RR",
        ],
        "files": ["train.txt", "valid.txt", "test.txt"],
    },
}


def download_dataset(name: str, out_dir: str = "data") -> str:
    if name not in SUPPORTED:
        raise ValueError(f"Unsupported dataset: {name}. Supported: {list(SUPPORTED)}")
    bases = SUPPORTED[name]["bases"]
    files = SUPPORTED[name]["files"]
    ds_dir = os.path.join(out_dir, name)
    os.makedirs(ds_dir, exist_ok=True)
    for fn in files:
        dest = os.path.join(ds_dir, fn)
        if os.path.exists(dest) and os.path.getsize(dest) > 0:
            continue
        last_err: Exception | None = None
        for base in bases:
            url = f"{base}/{fn}"
            try:
                r = requests.get(url, timeout=60)
                r.raise_for_status()
                with open(dest, "wb") as f:
                    f.write(r.content)
                last_err = None
                break
            except Exception as e:
                last_err = e
        if last_err is not None:
            raise last_err
    return ds_dir


def load_triples_from_dir(ds_dir: str, split: str = "train") -> List[Tuple[str, str, str]]:
    path = os.path.join(ds_dir, f"{split}.txt")
    if not os.path.exists(path):
        raise FileNotFoundError(path)
    triples: List[Tuple[str, str, str]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split()
            if len(parts) != 3:
                # Some datasets are tab-separated; try that
                parts = line.split("\t")
                if len(parts) != 3:
                    continue
            h, r, t = parts
            triples.append((h, r, t))
    return triples


def all_triples(ds_dir: str) -> List[Tuple[str, str, str]]:
    triples: List[Tuple[str, str, str]] = []
    for split in ("train", "valid", "test"):
        p = os.path.join(ds_dir, f"{split}.txt")
        if os.path.exists(p):
            triples.extend(load_triples_from_dir(ds_dir, split))
    return triples