Initial release: weights + ONNX + notebook

inference_examples.ipynb

@@ -0,0 +1,272 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "intro",
+   "metadata": {},
+   "source": [
+    "# programming-language-identification-100plus\n",
+    "\n",
+    "Runnable examples for the ModernBERT programming-language identifier.\n",
+    "Covers 107 languages. Input is truncated to the first 512 characters\n",
+    "(matches the training-time `head` strategy).\n",
+    "\n",
+    "Point `MODEL_ID` at the local checkpoint directory or the HF repo id."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "setup",
+   "metadata": {},
+   "outputs": [],
+   "source": "import torch\nfrom transformers import AutoModelForSequenceClassification, AutoTokenizer\n\nMODEL_ID = \"/home/vijay/llm_models/guardrail_code_models/programming-language-identification-100plus\"\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\ntokenizer = AutoTokenizer.from_pretrained(MODEL_ID)\nmodel = AutoModelForSequenceClassification.from_pretrained(\n    MODEL_ID,\n    attn_implementation=\"eager\",\n    torch_dtype=torch.bfloat16,  # weights are published in bf16\n).to(DEVICE).eval()\n\nprint(f\"device={DEVICE}  num_labels={model.config.num_labels}  dtype={model.dtype}\")\n"
+  },
+  {
+   "cell_type": "markdown",
+   "id": "helpers",
+   "metadata": {},
+   "source": [
+    "## Helpers"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "helpers-code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@torch.no_grad()\n",
+    "def predict(snippets, top_k=1, max_chars=512):\n",
+    "    \"\"\"Return the top-k languages + probabilities for each snippet.\"\"\"\n",
+    "    if isinstance(snippets, str):\n",
+    "        snippets = [snippets]\n",
+    "    trimmed = [s[:max_chars] for s in snippets]\n",
+    "    encoded = tokenizer(\n",
+    "        trimmed, return_tensors=\"pt\", padding=True, truncation=True, max_length=512\n",
+    "    ).to(DEVICE)\n",
+    "    logits = model(**encoded).logits\n",
+    "    probs = logits.softmax(-1)\n",
+    "    top_probs, top_ids = probs.topk(top_k, dim=-1)\n",
+    "    results = []\n",
+    "    for row_probs, row_ids in zip(top_probs.tolist(), top_ids.tolist()):\n",
+    "        results.append(\n",
+    "            [\n",
+    "                (model.config.id2label[label_id], prob)\n",
+    "                for label_id, prob in zip(row_ids, row_probs)\n",
+    "            ]\n",
+    "        )\n",
+    "    return results\n",
+    "\n",
+    "\n",
+    "def show(title, snippet, top_k=1):\n",
+    "    preds = predict(snippet, top_k=top_k)[0]\n",
+    "    head = snippet.strip().splitlines()[0][:60]\n",
+    "    print(f\"{title:14s}  `{head}`\")\n",
+    "    for name, prob in preds:\n",
+    "        print(f\"               {name:30s} {prob:.3f}\")\n",
+    "    print()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "single",
+   "metadata": {},
+   "source": [
+    "## 1. Single-snippet prediction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "single-code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "python_snippet = '''\n",
+    "def greet(name: str) -> None:\n",
+    "    print(f\"hello, {name}\")\n",
+    "\n",
+    "for person in [\"ada\", \"alan\", \"grace\"]:\n",
+    "    greet(person)\n",
+    "'''.strip()\n",
+    "\n",
+    "show(\"Python\", python_snippet)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "batch",
+   "metadata": {},
+   "source": [
+    "## 2. Batch across many languages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "batch-code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAMPLES = {\n",
+    "    \"Rust\": '''\n",
+    "fn main() {\n",
+    "    let names = vec![\"ada\", \"alan\", \"grace\"];\n",
+    "    for n in &names {\n",
+    "        println!(\"hello, {}\", n);\n",
+    "    }\n",
+    "}\n",
+    "'''.strip(),\n",
+    "    \"Go\": '''\n",
+    "package main\n",
+    "\n",
+    "import \"fmt\"\n",
+    "\n",
+    "func main() {\n",
+    "    names := []string{\"ada\", \"alan\", \"grace\"}\n",
+    "    for _, n := range names {\n",
+    "        fmt.Printf(\"hello, %s\\\\n\", n)\n",
+    "    }\n",
+    "}\n",
+    "'''.strip(),\n",
+    "    \"Ruby\": '''\n",
+    "[\"ada\", \"alan\", \"grace\"].each do |name|\n",
+    "  puts \"hello, #{name}\"\n",
+    "end\n",
+    "'''.strip(),\n",
+    "    \"Elixir\": '''\n",
+    "defmodule Greeter do\n",
+    "  def hello(name), do: IO.puts(\"hello, #{name}\")\n",
+    "end\n",
+    "\n",
+    "Enum.each([\"ada\", \"alan\", \"grace\"], &Greeter.hello/1)\n",
+    "'''.strip(),\n",
+    "    \"Haskell\": '''\n",
+    "main :: IO ()\n",
+    "main = mapM_ (\\\\n -> putStrLn (\"hello, \" ++ n)) [\"ada\", \"alan\", \"grace\"]\n",
+    "'''.strip(),\n",
+    "    \"Kotlin\": '''\n",
+    "fun main() {\n",
+    "    listOf(\"ada\", \"alan\", \"grace\").forEach { println(\"hello, $it\") }\n",
+    "}\n",
+    "'''.strip(),\n",
+    "    \"Mathematica/Wolfram Language\": '''\n",
+    "greet[name_String] := Print[\"hello, \" <> name];\n",
+    "greet /@ {\"ada\", \"alan\", \"grace\"};\n",
+    "'''.strip(),\n",
+    "    \"ARM Assembly\": '''\n",
+    "    .syntax unified\n",
+    "    .thumb\n",
+    "    .global main\n",
+    "main:\n",
+    "    ldr r0, =message\n",
+    "    bl  puts\n",
+    "    mov r0, #0\n",
+    "    bx  lr\n",
+    "message:\n",
+    "    .asciz \"hello\"\n",
+    "'''.strip(),\n",
+    "    \"Julia\": '''\n",
+    "for name in [\"ada\", \"alan\", \"grace\"]\n",
+    "    println(\"hello, $name\")\n",
+    "end\n",
+    "'''.strip(),\n",
+    "}\n",
+    "\n",
+    "snippets = list(SAMPLES.values())\n",
+    "expected = list(SAMPLES.keys())\n",
+    "predictions = predict(snippets, top_k=1)\n",
+    "\n",
+    "correct = 0\n",
+    "for gold, preds in zip(expected, predictions):\n",
+    "    predicted, prob = preds[0]\n",
+    "    mark = \"OK \" if predicted == gold else \"!  \"\n",
+    "    print(f\"  {mark} gold={gold:32s} pred={predicted:32s} p={prob:.3f}\")\n",
+    "    if predicted == gold:\n",
+    "        correct += 1\n",
+    "print(f\"\\n{correct}/{len(snippets)} correct\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "topk",
+   "metadata": {},
+   "source": [
+    "## 3. Top-k with confidence\n",
+    "\n",
+    "Useful when a snippet is short or ambiguous — inspect the runner-ups\n",
+    "before committing to a label."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "topk-code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Kotlin/Java syntactic overlap — see how far ahead the winner is\n",
+    "jvm_snippet = '''\n",
+    "class Hello {\n",
+    "    fun say(name: String) = println(\"hello, $name\")\n",
+    "}\n",
+    "'''.strip()\n",
+    "\n",
+    "show(\"JVM snippet\", jvm_snippet, top_k=5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ambiguous",
+   "metadata": {},
+   "source": [
+    "## 4. Very short / ambiguous input\n",
+    "\n",
+    "Snippets under ~60 characters are often genuinely ambiguous — multiple\n",
+    "languages accept the same syntax. Top-k probabilities will be diffuse."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ambiguous-code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "show(\"short\", \"x = 1\", top_k=5)\n",
+    "show(\"one-liner\", \"print('hi')\", top_k=5)\n",
+    "show(\"empty-ish\", \"{}\", top_k=5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "closing",
+   "metadata": {},
+   "source": [
+    "## Tips\n",
+    "\n",
+    "* Feed at least ~100 characters for reliable results.\n",
+    "* The model was trained and evaluated with the first 512 characters of each\n",
+    "  file. For longer files, that's also what you should pass.\n",
+    "* If you have file extensions available, treat them as a strong prior —\n",
+    "  this classifier is purely content-based and will happily misclassify a\n",
+    "  polyglot hello-world if you ask it to."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}