Load model from HuggingFace hub

README.md

@@ -0,0 +1,12 @@
+---
+title: ndml-eeg-code-generator
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 4.0.0
+app_file: app.py
+pinned: false
+---
+
+# OLMo EEG Code Generator
+NDML Lab — Parkinson's and EEG Analysis Code

app.py

@@ -0,0 +1,79 @@
+import os, torch, gradio as gr
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer
+from olmo_core.nn.transformer import TransformerConfig
+
+# Load from HuggingFace — works without lab server
+REPO_ID    = "Tbain20/olmo2-1b-eeg-v11"
+VOCAB_SIZE = 100278
+DEVICE     = "cpu"  # HF free tier is CPU only
+PREFIX_I   = "### Instruction:\n"
+PREFIX_R   = "\n\n### Response:\n"
+
+print("Downloading model from HuggingFace...")
+ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="best_model.pt")
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-2-1124-7B")
+
+print("Building model...")
+cfg   = TransformerConfig.olmo2_1B(vocab_size=VOCAB_SIZE)
+model = cfg.build()
+ckpt  = torch.load(ckpt_path, map_location="cpu")
+model.load_state_dict(ckpt["model_state_dict"])
+model.eval()
+print("Ready")
+
+@torch.no_grad()
+def generate_code(prompt, max_new_tokens=300, temperature=0.7):
+    if not prompt.strip():
+        return "Please enter a prompt."
+    full = PREFIX_I + prompt.strip() + PREFIX_R
+    ids  = tokenizer.encode(full)
+    x    = torch.tensor([ids], dtype=torch.long)
+    for _ in range(max_new_tokens):
+        logits     = model(x)
+        logits     = logits[:, -1, :] / temperature
+        vals, idxs = torch.topk(logits, k=40)
+        probs      = F.softmax(vals, dim=-1)
+        next_tok   = idxs.gather(-1, torch.multinomial(probs, 1))
+        x          = torch.cat([x, next_tok], dim=1)
+        if next_tok.item() == tokenizer.eos_token_id:
+            break
+    out = tokenizer.decode(x[0].tolist(), skip_special_tokens=True)
+    return out[len(full):].strip()
+
+EXAMPLES = [
+    ["Write a Python function using MNE to filter EEG data for beta waves (13-30 Hz)"],
+    ["Write a Python function to compute beta band power from STN LFP recordings"],
+    ["Write a Python function to compare beta power between on and off medication Parkinson's patients"],
+    ["Write a Python function to load TDT block and extract RSn1 LFP stream"],
+    ["Write a Python function to suppress DBS stimulation artifacts using sample-and-hold"],
+    ["Write a MATLAB function using FieldTrip to compute beta band power from LFP"],
+]
+
+with gr.Blocks(title="OLMo EEG Code Generator") as demo:
+    gr.Markdown("# 🧠 OLMo EEG Code Generator\n### NDML Lab — Parkinson's & EEG Analysis Assistant\n*Note: Running on CPU — generation takes 1-2 minutes*")
+    with gr.Row():
+        with gr.Column():
+            prompt_box  = gr.Textbox(label="Describe what you need", lines=4,
+                placeholder="e.g. Write a Python function using MNE to filter EEG for beta waves")
+            temperature = gr.Slider(0.3, 1.2, value=0.7, step=0.05, label="Temperature")
+            max_tokens  = gr.Slider(100, 500, value=300, step=50, label="Max tokens")
+            with gr.Row():
+                generate_btn = gr.Button("Generate Code", variant="primary", scale=2)
+                clear_btn    = gr.Button("Clear", scale=1)
+            gr.Examples(examples=EXAMPLES, inputs=prompt_box)
+        with gr.Column():
+            output_box = gr.Code(label="Generated Code", language="python", lines=25)
+
+    generate_btn.click(
+        fn=lambda p,t,m: generate_code(p, int(m), float(t)),
+        inputs=[prompt_box, temperature, max_tokens], outputs=output_box)
+    prompt_box.submit(
+        fn=lambda p,t,m: generate_code(p, int(m), float(t)),
+        inputs=[prompt_box, temperature, max_tokens], outputs=output_box)
+    clear_btn.click(fn=lambda: ("",""), outputs=[prompt_box, output_box])
+
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.44.0
+transformers
+torch
+numpy
+huggingface_hub
+ai2-olmo