vibe

README.md

@@ -9,4 +9,6 @@ app_file: app.py
 pinned: false
 ---
 
+This Space is set up for ZeroGPU via the `@spaces.GPU` decorator in `app.py`. Select ZeroGPU as the hardware in your Space settings.
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,124 @@
+import os
+import math
+import spaces
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL_ID = os.getenv("MODEL_ID", "Qwen/Qwen3-4B")
+
+
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    # Ensure a pad token exists for safe batching; use eos if needed
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model.eval()
+    return tokenizer, model
+
+
+TOKENIZER, MODEL = load_model()
+
+
+@spaces.GPU
+def compute_entropy(code: str):
+    if not code or not code.strip():
+        return "Please paste some source code.", None
+
+    with torch.no_grad():
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if next(MODEL.parameters()).device != device:
+            MODEL.to(device)
+
+        enc = TOKENIZER(code, return_tensors="pt")
+        input_ids = enc["input_ids"]
+        attention_mask = enc.get("attention_mask")
+
+        input_ids = input_ids.to(device)
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+
+        # Need at least 2 tokens to compute next-token NLL
+        if input_ids.shape[1] < 2:
+            return "Input is too short to compute token-level entropy.", None
+
+        outputs = MODEL(input_ids=input_ids, attention_mask=attention_mask)
+        logits = outputs.logits
+
+        # Shift for next-token prediction
+        shift_logits = logits[:, :-1, :]
+        shift_labels = input_ids[:, 1:]
+
+        log_probs = torch.log_softmax(shift_logits, dim=-1)
+        # Gather log prob of the true next token
+        true_log_probs = log_probs.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1)
+        nll = -true_log_probs  # negative log-likelihood (nats)
+
+        nll_list = nll.squeeze(0).detach().cpu().tolist()
+        label_ids = shift_labels.squeeze(0).detach().cpu().tolist()
+        tokens = TOKENIZER.convert_ids_to_tokens(label_ids)
+
+    rows = []
+    for tok, nll_val in zip(tokens, nll_list):
+        prob = math.exp(-nll_val)
+        rows.append([tok, float(nll_val), float(prob)])
+
+    avg_nll = sum(nll_list) / len(nll_list)
+    avg_bits = avg_nll / math.log(2)
+    summary = (
+        f"Tokens evaluated: {len(nll_list)}\n"
+        f"Average NLL (nats): {avg_nll:.4f}\n"
+        f"Average NLL (bits): {avg_bits:.4f}"
+    )
+
+    return summary, rows
+
+
+def build_app():
+    with gr.Blocks(title="Entropy for Source Code", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+# Source Code Entropy (Qwen3-4B)
+
+Paste code below to compute token-level negative log-likelihood (NLL).
+The table shows each token's NLL and probability under the model.
+"""
+        )
+
+        code = gr.Textbox(
+            label="Source Code",
+            lines=16,
+            placeholder="Paste your source code here...",
+        )
+        btn = gr.Button("Compute Entropy")
+        summary = gr.Textbox(label="Summary", lines=4)
+        table = gr.Dataframe(
+            headers=["token", "nll_nats", "prob"],
+            datatype=["str", "number", "number"],
+            label="Token-level NLL",
+        )
+
+        btn.click(fn=compute_entropy, inputs=[code], outputs=[summary, table])
+
+        gr.Markdown(
+            """
+Notes:
+- NLL is computed for next-token prediction and excludes the first token.
+- Large inputs may take time to process depending on hardware.
+"""
+        )
+
+    return demo
+
+
+app = build_app()
+
+if __name__ == "__main__":
+    app.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers>=4.45.0
+accelerate>=0.34.0
+torch>=2.2.0
+gradio>=6.5.1
+spaces>=0.28.0