Hf Enrich Models

1#!/usr/bin/env bash
2
3set -euo pipefail
4
5show_help() {
6    cat << 'USAGE'
7Stream model IDs on stdin, emit one JSON object per line (NDJSON).
8
9Usage:
10  hf_enrich_models.sh [MODEL_ID ...]
11  cat ids.txt | hf_enrich_models.sh
12  baseline_hf_api.sh 50 | jq -r '.[].id' | hf_enrich_models.sh
13
14Description:
15  Reads newline-separated model IDs and fetches basic metadata for each.
16  Outputs NDJSON with id, downloads, likes, pipeline_tag, tags.
17  Uses HF_TOKEN for auth if the environment variable is set.
18
19Examples:
20  hf_enrich_models.sh gpt2 distilbert-base-uncased
21  baseline_hf_api.sh 50 | jq -r '.[].id' | hf_enrich_models.sh | jq -s 'sort_by(.downloads)'
22  HF_TOKEN=your_token hf_enrich_models.sh microsoft/DialoGPT-medium
23USAGE
24}
25
26if [[ "${1:-}" == "--help" ]]; then
27    show_help
28    exit 0
29fi
30
31if ! command -v jq >/dev/null 2>&1; then
32    echo "Error: jq is required but not installed" >&2
33    exit 1
34fi
35
36headers=()
37if [[ -n "${HF_TOKEN:-}" ]]; then
38    headers=(-H "Authorization: Bearer ${HF_TOKEN}")
39fi
40
41emit_error() {
42    local model_id="$1"
43    local message="$2"
44    jq -cn --arg id "$model_id" --arg error "$message" '{id: $id, error: $error}'
45}
46
47process_id() {
48    local model_id="$1"
49
50    if [[ -z "$model_id" ]]; then
51        return 0
52    fi
53
54    local url="https://huggingface.co/api/models/${model_id}"
55    local response
56    response=$(curl -s "${headers[@]}" "$url" 2>/dev/null || true)
57
58    if [[ -z "$response" ]]; then
59        emit_error "$model_id" "request_failed"
60        return 0
61    fi
62
63    if ! jq -e . >/dev/null 2>&1 <<<"$response"; then
64        emit_error "$model_id" "invalid_json"
65        return 0
66    fi
67
68    if jq -e '.error' >/dev/null 2>&1 <<<"$response"; then
69        emit_error "$model_id" "not_found"
70        return 0
71    fi
72
73    jq -c --arg id "$model_id" '{
74        id: (.id // $id),
75        downloads: (.downloads // 0),
76        likes: (.likes // 0),
77        pipeline_tag: (.pipeline_tag // "unknown"),
78        tags: (.tags // [])
79    }' <<<"$response" 2>/dev/null || emit_error "$model_id" "parse_failed"
80}
81
82if [[ $# -gt 0 ]]; then
83    for model_id in "$@"; do
84        process_id "$model_id"
85    done
86    exit 0
87fi
88
89if [[ -t 0 ]]; then
90    show_help
91    exit 1
92fi
93
94while IFS= read -r model_id; do
95    process_id "$model_id"
96done