change api return and change logger

app.py

@@ -2,7 +2,7 @@ import collections
 import heapq
 import json
 import os
-import logging
+from accelerate.logging import get_logger
 
 import gradio as gr
 import numpy as np
@@ -13,9 +13,9 @@ from torchvision import transforms
 
 from templates import openai_imagenet_template
 
-log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
-logging.basicConfig(level=logging.INFO, format=log_format)
-logger = logging.getLogger()
+# log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
+# logging.basicConfig(level=logging.INFO, format=log_format)
+logger = get_logger(__name__, log_level="INFO")
 
 hf_token = os.getenv("HF_TOKEN")
 
@@ -155,7 +155,7 @@ def open_domain_classification(img, rank: int) -> dict[str, float]:
 
 
 @torch.no_grad()
-def api_classification(img, rank: int) -> dict[str, float]:
+def api_classification(img, rank: int): #  -> dict[str, float]:
     """
     Predicts from the entire tree of life.
     If targeting a higher rank than species, then this function predicts among all
@@ -182,10 +182,9 @@ def api_classification(img, rank: int) -> dict[str, float]:
 
     logger.info(">>>>")
     logger.info(probs[0])
-
-    topk_names = heapq.nlargest(k, output, key=output.get)
-
-    return {name: output[name] for name in topk_names}
+    return probs[0]
+    # topk_names = heapq.nlargest(k, output, key=output.get)
+    # return {name: output[name] for name in topk_names}
 
 
 def change_output(choice):