main.py

import logging
import os
import random
from datetime import timedelta
from statistics import mean
from typing import Annotated, Any, Iterator, Union

import fasttext
from cashews import cache
from dotenv import load_dotenv
from fastapi import FastAPI, Path, Query
from httpx import AsyncClient, Client, Timeout
from huggingface_hub import hf_hub_download
from iso639 import Lang
from starlette.responses import RedirectResponse
from toolz import concat, groupby, valmap

cache.setup("mem://")


logger = logging.getLogger(__name__)
app = FastAPI()
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
assert HF_TOKEN
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

FASTTEXT_PREFIX_LENGTH = 9  # fasttext labels are formatted like "__label__eng_Latn"

BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
DEFAULT_FAST_TEXT_MODEL = "facebook/fasttext-language-identification"
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

timeout = Timeout(60, read=120)
client = Client(headers=headers, timeout=timeout)
async_client = AsyncClient(headers=headers, timeout=timeout)

TARGET_COLUMN_NAMES = {
    "text",
    "input",
    "tokens",
    "prompt",
    "instruction",
    "sentence_1",
    "question",
    "sentence2",
    "answer",
    "sentence",
    "response",
    "context",
    "query",
    "chosen",
    "rejected",
    "question"
}


def datasets_server_valid_rows(hub_id: str):
    try:
        resp = client.get(f"{BASE_DATASETS_SERVER_URL}/is-valid?dataset={hub_id}")
        data = resp.json()
        return True if data.get("viewer") else bool(data.get("preview"))
    except Exception as e:
        logger.error(f"Failed to get is-valid for {hub_id}: {e}")
        return False


async def get_first_config_and_split_name(hub_id: str):
    try:
        resp = await async_client.get(
            f"https://datasets-server.huggingface.co/splits?dataset={hub_id}"
        )

        data = resp.json()
        return data["splits"][0]["config"], data["splits"][0]["split"]
    except Exception as e:
        logger.error(f"Failed to get splits for {hub_id}: {e}")
        return (None, None)  # Return a tuple of None values


async def get_dataset_info(hub_id: str, config: str | None = None):
    if config is None:
        config_tuple, _ = await get_first_config_and_split_name(hub_id)
        if config_tuple is None:
            return None
        else:
            config = config_tuple
    resp = await async_client.get(
        f"{BASE_DATASETS_SERVER_URL}/info?dataset={hub_id}&config={config}"
    )
    resp.raise_for_status()
    return resp.json()


@cache(ttl=timedelta(minutes=5))
async def fetch_rows(url: str) -> list[dict]:
    response = await async_client.get(url)
    if response.status_code == 200:
        data = response.json()
        return data.get("rows")
    else:
        print(f"Failed to fetch data: {response.status_code}")
        print(url)
        return []


# Function to get random rows from the dataset
async def get_random_rows(
    hub_id: str,
    total_length: int,
    number_of_rows: int,
    max_request_calls: int,
    config="default",
    split="train",
):
    rows = []
    rows_per_call = min(
        number_of_rows // max_request_calls, total_length // max_request_calls
    )
    rows_per_call = min(rows_per_call, 100)  # Ensure rows_per_call is not more than 100
    for _ in range(min(max_request_calls, number_of_rows // rows_per_call)):
        offset = random.randint(0, total_length - rows_per_call)
        url = f"https://datasets-server.huggingface.co/rows?dataset={hub_id}&config={config}&split={split}&offset={offset}&length={rows_per_call}"
        logger.info(f"Fetching {url}")
        batch_rows = await fetch_rows(url)
        rows.extend(batch_rows)
        if len(rows) >= number_of_rows:
            break
    return [row.get("row") for row in rows]


def load_model(repo_id: str) -> fasttext.FastText._FastText:
    from pathlib import Path

    Path("code/models").mkdir(parents=True, exist_ok=True)
    model_path = hf_hub_download(
        repo_id,
        "model.bin",
        # cache_dir="code/models",
        # local_dir="code/models",
        # local_dir_use_symlinks=False,
    )
    return fasttext.load_model(model_path)


model = load_model(DEFAULT_FAST_TEXT_MODEL)


def yield_clean_rows(rows: Union[list[str], str], min_length: int = 3) -> Iterator[str]:
    for row in rows:
        if isinstance(row, str):
            # split on lines and remove empty lines
            line = row.split("\n")
            for line in line:
                if line:
                    yield line
        elif isinstance(row, list):
            try:
                line = " ".join(row)
                if len(line) < min_length:
                    continue
                else:
                    yield line
            except TypeError:
                continue


def model_predict(inputs: str, k=1) -> list[dict[str, float]]:
    predictions = model.predict(inputs, k=k)
    return [
        {"label": label[FASTTEXT_PREFIX_LENGTH:], "score": prob}
        for label, prob in zip(predictions[0], predictions[1])
    ]


def get_label(x):
    return x.get("label")


def get_mean_score(preds):
    return mean([pred.get("score") for pred in preds])


def filter_by_frequency(counts_dict: dict, threshold_percent: float = 0.2):
    """Filter a dict to include items whose value is above `threshold_percent`"""
    total = sum(counts_dict.values())
    threshold = total * threshold_percent
    return {k for k, v in counts_dict.items() if v >= threshold}


def try_parse_language(lang: str) -> str | None:
    try:
        split = lang.split("_")
        lang = split[0]
        lang = Lang(lang)
        return lang.pt1
    except Exception as e:
        logger.error(f"Failed to parse language {lang}: {e}")
        return None


def predict_rows(
    rows, target_column, language_threshold_percent=0.2, return_raw_predictions=False
):
    rows = (row.get(target_column) for row in rows)
    rows = (row for row in rows if row is not None)
    rows = list(yield_clean_rows(rows))
    predictions = [model_predict(row) for row in rows]
    predictions = [pred for pred in predictions if pred is not None]
    predictions = list(concat(predictions))
    predictions_by_lang = groupby(get_label, predictions)
    langues_counts = valmap(len, predictions_by_lang)
    keys_to_keep = filter_by_frequency(
        langues_counts, threshold_percent=language_threshold_percent
    )
    filtered_dict = {k: v for k, v in predictions_by_lang.items() if k in keys_to_keep}
    raw_model_prediction_summary = dict(valmap(get_mean_score, filtered_dict))
    parsed_langs = {
        try_parse_language(k): v for k, v in raw_model_prediction_summary.items()
    }
    default_data = {
        "language_prediction_summary": parsed_langs,
        "raw_model_prediction_summary": raw_model_prediction_summary,
        "hub_id": "hub_id",
        "config": "config",
    }
    if return_raw_predictions:
        default_data["raw_predictions"] = predictions
    return default_data


@app.get("/", include_in_schema=False)
def root():
    return RedirectResponse(url="/docs")


@app.get("/predict_dataset_language/{hub_id:path}")
@cache(ttl=timedelta(minutes=10))
async def predict_language(
    hub_id: Annotated[str, Path(title="The hub id of the dataset to predict")],
    config: str | None = None,
    split: str | None = None,
    max_request_calls: Annotated[
        int, Query(title="Max number of requests to datasets server", gt=0, le=50)
    ] = 10,
    number_of_rows: int = 1000,
    language_threshold_percent: float = 0.2,
) -> dict[Any, Any] | None:
    is_valid = datasets_server_valid_rows(hub_id)
    if not is_valid:
        logger.error(f"Dataset {hub_id} is not accessible via the datasets server.")
        return None  # Return early if dataset is not valid
        
    if not config and not split:
        config_tuple, split_tuple = await get_first_config_and_split_name(hub_id)
        if config_tuple is None:
            logger.error(f"Could not retrieve configuration for dataset {hub_id}")
            return None
        config, split = config_tuple, split_tuple
    elif not config:
        config_tuple, _ = await get_first_config_and_split_name(hub_id)
        if config_tuple is None:
            logger.error(f"Could not retrieve configuration for dataset {hub_id}")
            return None
        config = config_tuple
    elif not split:
        _, split_tuple = await get_first_config_and_split_name(hub_id)
        if split_tuple is None:
            logger.error(f"Could not retrieve split for dataset {hub_id}")
            return None
        split = split_tuple
        
    info = await get_dataset_info(hub_id, config)
    if info is None:
        logger.error(f"Dataset {hub_id} is not accessible via the datasets server.")
        return None
        
    if dataset_info := info.get("dataset_info"):
        total_rows_for_split = dataset_info.get("splits").get(split).get("num_examples")
        features = dataset_info.get("features")
        
        # Get original column names
        column_names = set(features.keys())
        logger.info(f"Column names: {column_names}")
        
        # Create a mapping of lowercase column names to their original casing
        lowercase_to_original = {col.lower(): col for col in column_names}
        
        # Check intersection with lowercase versions
        lowercase_column_names = set(lowercase_to_original.keys())
        lowercase_target_columns = {col.lower() for col in TARGET_COLUMN_NAMES}
        
        if not lowercase_column_names.intersection(lowercase_target_columns):
            logger.error(
                f"Dataset {hub_id} {column_names} does not contain any of the target columns {TARGET_COLUMN_NAMES}"
            )
            return None
            
        # Find target column with case-insensitive matching
        target_column = None
        for column in TARGET_COLUMN_NAMES:
            if column.lower() in lowercase_column_names:
                # Use the original casing from the dataset
                target_column = lowercase_to_original[column.lower()]
                logger.info(f"Using column {target_column} for language detection")
                break
                
        if target_column is None:
            logger.error(f"Could not find a suitable column for language detection")
            return None
            
        random_rows = await get_random_rows(
            hub_id,
            total_rows_for_split,
            number_of_rows,
            max_request_calls,
            config,
            split,
        )
        
        logger.info(f"Predicting language for {len(random_rows)} rows")
        predictions = predict_rows(
            random_rows,
            target_column,
            language_threshold_percent=language_threshold_percent,
        )
        predictions["hub_id"] = hub_id
        predictions["config"] = config
        predictions["split"] = split
        return predictions
        
    else:
        logger.error(f"No dataset_info available for {hub_id}")
        return None