Use transformers AutoModelForCausalLM — no OLMo-Core dependency

app.py

@@ -1,79 +1,74 @@
-import os, torch, gradio as gr
+import os
+import torch
+import 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
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# 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"])
+print("Loading model from HuggingFace...")
+model = AutoModelForCausalLM.from_pretrained(
+    REPO_ID,
+    torch_dtype=torch.float32,
+    device_map="cpu",
+)
 model.eval()
-print("Ready")
+print("Model 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()
+    ids  = tokenizer(full, return_tensors="pt").input_ids
+    out  = model.generate(
+        ids,
+        max_new_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+        do_sample=True,
+        top_k=40,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    generated = out[0][ids.shape[1]:]
+    return tokenizer.decode(generated, skip_special_tokens=True).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 compute beta band power from STN LFP using Welch method"],
     ["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 load TDT block and extract RSn1 LFP stream at 24414 Hz"],
     ["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*")
+    gr.Markdown("""# 🧠 OLMo EEG Code Generator
+### NDML Lab — Parkinson's & EEG Analysis Assistant
+*Neural Dynamics and Modulation Lab, Cleveland Clinic*
+
+> ⚠️ Running on CPU — generation takes 2-3 minutes. For fast generation use the lab server demo.""")
     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():
+                temperature = gr.Slider(0.3, 1.2, value=0.7, step=0.05, label="Temperature")
+                max_tokens  = gr.Slider(100, 400, 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)
+            gr.Examples(examples=EXAMPLES, inputs=prompt_box, label="Example prompts")
         with gr.Column():
-            output_box = gr.Code(label="Generated Code", language="python", lines=25)
+            output_box = gr.Code(label="Generated Code", language="python", lines=22)
 
-    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])
+    generate_btn.click(fn=generate_code, inputs=[prompt_box, temperature, max_tokens], outputs=output_box)
+    prompt_box.submit(fn=generate_code, inputs=[prompt_box, temperature, max_tokens], outputs=output_box)
+    clear_btn.click(fn=lambda: ("", ""), outputs=[prompt_box, output_box])
 
 demo.launch()

requirements.txt

@@ -1,6 +1,5 @@
 gradio==4.44.0
-transformers
-torch
-numpy
+torch==2.1.0
+transformers>=4.40.0
 huggingface_hub
-ai2-olmo
+numpy