| 1 | module ModelsHelper |
| 2 | # A small curated map of Hugging Face pipeline tags to human labels. Unknown |
| 3 | # tags fall back to a titleized version of the raw value. |
| 4 | PIPELINE_LABELS = { |
| 5 | "text-generation" => "Text Generation", |
| 6 | "text2text-generation" => "Text-to-Text", |
| 7 | "text-classification" => "Text Classification", |
| 8 | "token-classification" => "Token Classification", |
| 9 | "question-answering" => "Question Answering", |
| 10 | "summarization" => "Summarization", |
| 11 | "translation" => "Translation", |
| 12 | "feature-extraction" => "Feature Extraction", |
| 13 | "sentence-similarity" => "Sentence Similarity", |
| 14 | "fill-mask" => "Fill-Mask", |
| 15 | "image-classification" => "Image Classification", |
| 16 | "image-to-text" => "Image-to-Text", |
| 17 | "text-to-image" => "Text-to-Image", |
| 18 | "automatic-speech-recognition" => "Speech Recognition", |
| 19 | "text-to-speech" => "Text-to-Speech", |
| 20 | "audio-classification" => "Audio Classification" |
| 21 | }.freeze |
| 22 | |
| 23 | def pipeline_label(tag) |
| 24 | return nil if tag.blank? |
| 25 | PIPELINE_LABELS[tag] || tag.tr("-_", " ").split.map(&:capitalize).join(" ") |
| 26 | end |
| 27 | |
| 28 | # Compact human download count, e.g. 1240 → "1.2k", 3_400_000 → "3.4M". |
| 29 | def humanize_count(n) |
| 30 | n = n.to_i |
| 31 | return n.to_s if n < 1000 |
| 32 | number_to_human(n, units: { thousand: "k", million: "M", billion: "B" }, |
| 33 | format: "%n%u", precision: 2, strip_insignificant_zeros: true) |
| 34 | end |
| 35 | |
| 36 | # Maps a weight-file extension to a short format badge label. |
| 37 | WEIGHT_FORMATS = { |
| 38 | "safetensors" => "Safetensors", |
| 39 | "gguf" => "GGUF", |
| 40 | "bin" => "PyTorch", |
| 41 | "pt" => "PyTorch", |
| 42 | "pth" => "PyTorch", |
| 43 | "onnx" => "ONNX", |
| 44 | "ckpt" => "Checkpoint", |
| 45 | "h5" => "Keras", |
| 46 | "msgpack" => "Flax", |
| 47 | "tflite" => "TFLite" |
| 48 | }.freeze |
| 49 | |
| 50 | def weight_format(filename) |
| 51 | WEIGHT_FORMATS[File.extname(filename).delete_prefix(".").downcase] |
| 52 | end |
| 53 | end |