From d675b4ae18142ad7f20cb97f09d0fc4c9bfc4db6 Mon Sep 17 00:00:00 2001
From: chand1012 <chandler@chand1012.dev>
Date: Thu, 15 Jul 2021 21:56:17 +0000
Subject: [PATCH] Add timer

---
 api.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/api.py b/api.py
index 12b05a323..46599b9c2 100644
--- a/api.py
+++ b/api.py
@@ -1,4 +1,5 @@
 from typing import Optional
+import time
 
 from fastapi import FastAPI
 
@@ -24,9 +25,13 @@ def samples_get(
     temperature: Optional[float] = 1.0
 ):
 
+    start = time.time() * 1000
     output = generate.samples(prompt, model_name, seed, nsamples, batch_size, length, temperature, top_k)
+    end = time.time() * 1000
 
-    return {'prompt':prompt, 'responses':output}
+    diff = end - start
+
+    return {'prompt': prompt, 'responses': output, 'time': diff}
 
 # @app.post('/samples', response_class=ResponseModel)
 # def samples_post(data: RequestModel):