refactor: output structure

app.py

@@ -8,8 +8,13 @@ import spaces
 import logging
 import datetime
 
+# Load model for first time cache
+model = Detoxify("unbiased-small")
+
+
 @spaces.GPU
 def classify(query):
+    torch_device = "cuda" if torch.cuda.is_available() else "cpu"
     model = Detoxify("unbiased-small", device="cuda")
 
     all_result = []
@@ -25,27 +30,31 @@ def classify(query):
         data = [query]
         pass
 
+    start_time = datetime.datetime.now()
     for i in range(len(data)):
         result = {}
-        start_time = datetime.datetime.now()
-        
+
         df = pd.DataFrame(model.predict(str(data[i])), index=[0])
         columns = df.columns
 
         for i, label in enumerate(columns):
             result[label] = df[label][0].round(3).astype("float")
 
-        end_time = datetime.datetime.now()
-        elapsed_time = end_time - start_time
-        result["time"] =  str(elapsed_time)
-
-        logging.debug("elapsed predict time: %s", str(elapsed_time))
-        print("elapsed predict time:", str(elapsed_time))
-        
         all_result.append(result)
 
-    
-    return json.dumps(all_result) if request_type == list else all_result[0]
+    end_time = datetime.datetime.now()
+    elapsed_time = end_time - start_time
+
+    logging.debug("elapsed predict time: %s", str(elapsed_time))
+    print("elapsed predict time:", str(elapsed_time))
+
+    output = {}
+    output["time"] = str(elapsed_time)
+    output["device"] = torch_device
+    output["result"] = all_result
+
+    return json.dumps(output)
+
 
 demo = gr.Interface(fn=classify, inputs=["text"], outputs="text")
 demo.launch()