wip

demo/python/picollm_demo_completion.py

@@ -8,7 +8,7 @@
 # an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations under the License.
 #
-
+import time
 from argparse import ArgumentParser
 
 import picollm
@@ -131,6 +131,14 @@ def main():
     print(f"picoLLM {o.version}")
     print(f"Loaded {o.model}")
 
+    start_sec = [0.]
+
+    def stream_callback(x: str) -> None:
+        if start_sec[0] == 0.:
+            start_sec[0] = time.time()
+
+        print(x, flush=True, end='')
+
     try:
         completion = o.generate(
             prompt=prompt,
@@ -142,11 +150,14 @@ def main():
             temperature=temperature,
             top_p=top_p,
             num_top_choices=num_top_choices,
-            stream_callback=lambda x: print(x, flush=True, end=''))
+            stream_callback=stream_callback)
         print('\n')
 
         if verbose:
             print(completion)
+
+        tos = (completion.usage.completion_tokens - 1) / (time.time() - start_sec[0])
+        print(f"Generated {tos:.2f} tokens per second")
     finally:
         o.release()