From a01b48814ae5918295efab348cd89e55d780ae27 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 17 Jan 2024 14:37:07 -0800
Subject: [PATCH] add example for pulling with progress bar

---
 examples/pull-stream-progress/main.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 examples/pull-stream-progress/main.py

diff --git a/examples/pull-stream-progress/main.py b/examples/pull-stream-progress/main.py
new file mode 100644
index 0000000..02db1af
--- /dev/null
+++ b/examples/pull-stream-progress/main.py
@@ -0,0 +1,23 @@
+# pip install tqdm
+
+from tqdm import tqdm
+from ollama import pull
+
+
+current_digest, bars = '', {}
+for progress in pull('mistral', stream=True):
+  digest = progress.get('digest', '')
+  if digest != current_digest and current_digest in bars:
+    bars[current_digest].close()
+
+  if not digest:
+    print(progress.get('status'))
+    continue
+
+  if digest not in bars and (total := progress.get('total')):
+    bars[digest] = tqdm(total=total, desc=f'pushing {digest[7:19]}', unit='B', unit_scale=True)
+
+  if completed := progress.get('completed'):
+    bars[digest].update(completed - bars[digest].n)
+
+  current_digest = digest