librarian-bots/dataset-language-detection-api

1import logging2import os3import random4from datetime import timedelta5from statistics import mean6from typing import Annotated, Any, Iterator, Union7 8import fasttext9from cashews import cache10from dotenv import load_dotenv11from fastapi import FastAPI, Path, Query12from httpx import AsyncClient, Client, Timeout13from huggingface_hub import hf_hub_download14from iso639 import Lang15from starlette.responses import RedirectResponse16from toolz import concat, groupby, valmap17 18cache.setup("mem://")19 20 21logger = logging.getLogger(__name__)22app = FastAPI()23load_dotenv()24HF_TOKEN = os.getenv("HF_TOKEN")25assert HF_TOKEN26os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"27 28FASTTEXT_PREFIX_LENGTH = 9  # fasttext labels are formatted like "__label__eng_Latn"29 30BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"31DEFAULT_FAST_TEXT_MODEL = "facebook/fasttext-language-identification"32headers = {"Authorization": f"Bearer {HF_TOKEN}"}33 34timeout = Timeout(60, read=120)35client = Client(headers=headers, timeout=timeout)36async_client = AsyncClient(headers=headers, timeout=timeout)37 38TARGET_COLUMN_NAMES = {39    "text",40    "input",41    "tokens",42    "prompt",43    "instruction",44    "sentence_1",45    "question",46    "sentence2",47    "answer",48    "sentence",49    "response",50    "context",51    "query",52    "chosen",53    "rejected",54    "question"55}56 57 58def datasets_server_valid_rows(hub_id: str):59    try:60        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={hub_id}")61        data = resp.json()62        return True if data.get("viewer") else bool(data.get("preview"))63    except Exception as e:64        logger.error(f"Failed to get is-valid for {hub_id}: {e}")65        return False66 67 68async def get_first_config_and_split_name(hub_id: str):69    try:70        resp = await async_client.get(71            f"https://datasets-server.huggingface.co/splits?dataset={hub_id}"72        )73 74        data = resp.json()75        return data["splits"][0]["config"], data["splits"][0]["split"]76    except Exception as e:77        logger.error(f"Failed to get splits for {hub_id}: {e}")78        return (None, None)  # Return a tuple of None values79 80 81async def get_dataset_info(hub_id: str, config: str | None = None):82    if config is None:83        config_tuple, _ = await get_first_config_and_split_name(hub_id)84        if config_tuple is None:85            return None86        else:87            config = config_tuple88    resp = await async_client.get(89        f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}"90    )91    resp.raise_for_status()92    return resp.json()93 94 95@cache(ttl=timedelta(minutes=5))96async def fetch_rows(url: str) -> list[dict]:97    response = await async_client.get(url)98    if response.status_code == 200:99        data = response.json()100        return data.get("rows")101    else:102        print(f"Failed to fetch data: {response.status_code}")103        print(url)104        return []105 106 107# Function to get random rows from the dataset108async def get_random_rows(109    hub_id: str,110    total_length: int,111    number_of_rows: int,112    max_request_calls: int,113    config="default",114    split="train",115):116    rows = []117    rows_per_call = min(118        number_of_rows // max_request_calls, total_length // max_request_calls119    )120    rows_per_call = min(rows_per_call, 100)  # Ensure rows_per_call is not more than 100121    for _ in range(min(max_request_calls, number_of_rows // rows_per_call)):122        offset = random.randint(0, total_length - rows_per_call)123        url = f"https://datasets-server.huggingface.co/rows?dataset={hub_id}&config={config}&split={split}&offset={offset}&length={rows_per_call}"124        logger.info(f"Fetching {url}")125        batch_rows = await fetch_rows(url)126        rows.extend(batch_rows)127        if len(rows) >= number_of_rows:128            break129    return [row.get("row") for row in rows]130 131 132def load_model(repo_id: str) -> fasttext.FastText._FastText:133    from pathlib import Path134 135    Path("code/models").mkdir(parents=True, exist_ok=True)136    model_path = hf_hub_download(137        repo_id,138        "model.bin",139        # cache_dir="code/models",140        # local_dir="code/models",141        # local_dir_use_symlinks=False,142    )143    return fasttext.load_model(model_path)144 145 146model = load_model(DEFAULT_FAST_TEXT_MODEL)147 148 149def yield_clean_rows(rows: Union[list[str], str], min_length: int = 3) -> Iterator[str]:150    for row in rows:151        if isinstance(row, str):152            # split on lines and remove empty lines153            line = row.split("\n")154            for line in line:155                if line:156                    yield line157        elif isinstance(row, list):158            try:159                line = " ".join(row)160                if len(line) < min_length:161                    continue162                else:163                    yield line164            except TypeError:165                continue166 167 168def model_predict(inputs: str, k=1) -> list[dict[str, float]]:169    predictions = model.predict(inputs, k=k)170    return [171        {"label": label[FASTTEXT_PREFIX_LENGTH:], "score": prob}172        for label, prob in zip(predictions[0], predictions[1])173    ]174 175 176def get_label(x):177    return x.get("label")178 179 180def get_mean_score(preds):181    return mean([pred.get("score") for pred in preds])182 183 184def filter_by_frequency(counts_dict: dict, threshold_percent: float = 0.2):185    """Filter a dict to include items whose value is above `threshold_percent`"""186    total = sum(counts_dict.values())187    threshold = total * threshold_percent188    return {k for k, v in counts_dict.items() if v >= threshold}189 190 191def try_parse_language(lang: str) -> str | None:192    try:193        split = lang.split("_")194        lang = split[0]195        lang = Lang(lang)196        return lang.pt1197    except Exception as e:198        logger.error(f"Failed to parse language {lang}: {e}")199        return None200 201 202def predict_rows(203    rows, target_column, language_threshold_percent=0.2, return_raw_predictions=False204):205    rows = (row.get(target_column) for row in rows)206    rows = (row for row in rows if row is not None)207    rows = list(yield_clean_rows(rows))208    predictions = [model_predict(row) for row in rows]209    predictions = [pred for pred in predictions if pred is not None]210    predictions = list(concat(predictions))211    predictions_by_lang = groupby(get_label, predictions)212    langues_counts = valmap(len, predictions_by_lang)213    keys_to_keep = filter_by_frequency(214        langues_counts, threshold_percent=language_threshold_percent215    )216    filtered_dict = {k: v for k, v in predictions_by_lang.items() if k in keys_to_keep}217    raw_model_prediction_summary = dict(valmap(get_mean_score, filtered_dict))218    parsed_langs = {219        try_parse_language(k): v for k, v in raw_model_prediction_summary.items()220    }221    default_data = {222        "language_prediction_summary": parsed_langs,223        "raw_model_prediction_summary": raw_model_prediction_summary,224        "hub_id": "hub_id",225        "config": "config",226    }227    if return_raw_predictions:228        default_data["raw_predictions"] = predictions229    return default_data230 231 232@app.get("/", include_in_schema=False)233def root():234    return RedirectResponse(url="/docs")235 236 237@app.get("/predict_dataset_language/{hub_id:path}")238@cache(ttl=timedelta(minutes=10))239async def predict_language(240    hub_id: Annotated[str, Path(title="The hub id of the dataset to predict")],241    config: str | None = None,242    split: str | None = None,243    max_request_calls: Annotated[244        int, Query(title="Max number of requests to datasets server", gt=0, le=50)245    ] = 10,246    number_of_rows: int = 1000,247    language_threshold_percent: float = 0.2,248) -> dict[Any, Any] | None:249    is_valid = datasets_server_valid_rows(hub_id)250    if not is_valid:251        logger.error(f"Dataset {hub_id} is not accessible via the datasets server.")252        return None  # Return early if dataset is not valid253        254    if not config and not split:255        config_tuple, split_tuple = await get_first_config_and_split_name(hub_id)256        if config_tuple is None:257            logger.error(f"Could not retrieve configuration for dataset {hub_id}")258            return None259        config, split = config_tuple, split_tuple260    elif not config:261        config_tuple, _ = await get_first_config_and_split_name(hub_id)262        if config_tuple is None:263            logger.error(f"Could not retrieve configuration for dataset {hub_id}")264            return None265        config = config_tuple266    elif not split:267        _, split_tuple = await get_first_config_and_split_name(hub_id)268        if split_tuple is None:269            logger.error(f"Could not retrieve split for dataset {hub_id}")270            return None271        split = split_tuple272        273    info = await get_dataset_info(hub_id, config)274    if info is None:275        logger.error(f"Dataset {hub_id} is not accessible via the datasets server.")276        return None277        278    if dataset_info := info.get("dataset_info"):279        total_rows_for_split = dataset_info.get("splits").get(split).get("num_examples")280        features = dataset_info.get("features")281        282        # Get original column names283        column_names = set(features.keys())284        logger.info(f"Column names: {column_names}")285        286        # Create a mapping of lowercase column names to their original casing287        lowercase_to_original = {col.lower(): col for col in column_names}288        289        # Check intersection with lowercase versions290        lowercase_column_names = set(lowercase_to_original.keys())291        lowercase_target_columns = {col.lower() for col in TARGET_COLUMN_NAMES}292        293        if not lowercase_column_names.intersection(lowercase_target_columns):294            logger.error(295                f"Dataset {hub_id} {column_names} does not contain any of the target columns {TARGET_COLUMN_NAMES}"296            )297            return None298            299        # Find target column with case-insensitive matching300        target_column = None301        for column in TARGET_COLUMN_NAMES:302            if column.lower() in lowercase_column_names:303                # Use the original casing from the dataset304                target_column = lowercase_to_original[column.lower()]305                logger.info(f"Using column {target_column} for language detection")306                break307                308        if target_column is None:309            logger.error(f"Could not find a suitable column for language detection")310            return None311            312        random_rows = await get_random_rows(313            hub_id,314            total_rows_for_split,315            number_of_rows,316            max_request_calls,317            config,318            split,319        )320        321        logger.info(f"Predicting language for {len(random_rows)} rows")322        predictions = predict_rows(323            random_rows,324            target_column,325            language_threshold_percent=language_threshold_percent,326        )327        predictions["hub_id"] = hub_id328        predictions["config"] = config329        predictions["split"] = split330        return predictions331        332    else:333        logger.error(f"No dataset_info available for {hub_id}")334        return None