From 397826f0b44737059073d8a5c4068cf70675cb18 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Mon, 7 Jul 2025 18:47:25 -0700
Subject: [PATCH] add a test for 1B llm (#11124)

* add a test for 1B llm

* fix mbs

* add apps to release
---
 .github/workflows/test.yml | 14 ++++++++++++++
 setup.py                   |  2 +-
 tinygrad/apps/llm.py       |  5 ++++-
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 6dd1a26bda..3c7c8ab549 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -547,6 +547,20 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
+  testllm:
+    name: Test LLM
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: apps_llm
+      - name: Test 1B LLM
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
+
   testmodels:
     name: Models (llvm+cpu+gpu)
     runs-on: ubuntu-22.04
diff --git a/setup.py b/setup.py
index 5d4d0fc3c6..a30c8fe56f 100644
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,7 @@ setup(name='tinygrad',
       packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
                   'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize',
                   'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv'],
+                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
       package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py
index 001790920e..baffe85aeb 100644
--- a/tinygrad/apps/llm.py
+++ b/tinygrad/apps/llm.py
@@ -172,7 +172,10 @@ if __name__ == "__main__":
   ids: list[int] = [bos_id]
   while 1:
     start_pos = len(ids) - 1
-    ids += tok.role("user") + tok.encode(input('>>> ')) + [eos_id] + tok.role("assistant")
+    try:
+      ids += tok.role("user") + tok.encode(input('>>> ')) + [eos_id] + tok.role("assistant")
+    except EOFError:
+      break
     for next_id in model.generate(ids, start_pos):
       sys.stdout.write(tok.decode([next_id]) if next_id != eos_id else "\n\n")
       sys.stdout.flush()