Create app.py

app.py

@@ -0,0 +1,95 @@
+import torch
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_ID = "UniversalComputingResearch/Atom2.7m"
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+).eval()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+
+def generate(prompt, max_new_tokens, temperature, do_sample):
+    if not prompt.strip():
+        return "Enter a prompt first."
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        add_special_tokens=False,
+    ).to(device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=bool(do_sample),
+            temperature=float(temperature) if do_sample else None,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+
+examples = [
+    ["12 + 34 =", 4, 1.0, False],
+    ["7 + 8 =", 3, 1.0, False],
+    ["25 - 9 =", 4, 1.0, False],
+    ["3 * 6 =", 4, 1.0, False],
+    ["The capital of France is", 12, 0.8, True],
+]
+
+description = """
+Atom2.7m is a tiny causal language model for text continuation, with arithmetic-aware handling for numeric spans.
+
+It is not an instruction-tuned chatbot. It works best with short continuation prompts such as `12 + 34 =`.
+"""
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Textbox(
+            label="Prompt",
+            value="12 + 34 =",
+            lines=3,
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=64,
+            value=8,
+            step=1,
+            label="Max new tokens",
+        ),
+        gr.Slider(
+            minimum=0.1,
+            maximum=2.0,
+            value=1.0,
+            step=0.1,
+            label="Temperature",
+        ),
+        gr.Checkbox(
+            value=False,
+            label="Sample instead of greedy decoding",
+        ),
+    ],
+    outputs=gr.Textbox(
+        label="Model output",
+        lines=6,
+    ),
+    title="Atom2.7m Arithmetic Demo",
+    description=description,
+    examples=examples,
+    allow_flagging="never",
+)
+
+if __name__ == "__main__":
+    demo.launch()