Add GSM8K

app.py

@@ -89,6 +89,17 @@ def init_space():
         EVAL_REQUESTS_PATH, EVAL_COLS
     )
     return dataset_df, original_df, finished_eval_queue_df, running_eval_queue_df, pending_eval_queue_df
+    
+    
+def add_benchmark_columns(shown_columns):
+    benchmark_columns = []
+    for benchmark in BENCHMARK_COLS:
+        if benchmark in shown_columns:
+            for c in COLS:
+                if benchmark in c and benchmark != c:
+                    benchmark_columns.append(c)
+    return benchmark_columns
+
 
 # Searching and filtering
 def update_table(
@@ -96,7 +107,8 @@ def update_table(
 ):
     filtered_df = filter_models(hidden_df, type_query, size_query, precision_query)
     filtered_df = filter_queries(query, filtered_df)
-    df = select_columns(filtered_df, columns)
+    benchmark_columns = add_benchmark_columns(columns)
+    df = select_columns(filtered_df, columns + benchmark_columns)
     return df
 
 
@@ -270,18 +282,19 @@ with demo:
                     # )
 
             # breakpoint()
-
+            benchmark_columns = add_benchmark_columns(shown_columns.value)
             leaderboard_table = gr.components.Dataframe(
                 value=(
                     leaderboard_df[
                         [c.name for c in fields(AutoEvalColumn) if c.never_hidden]
                         + shown_columns.value
+                        + benchmark_columns
                         + [AutoEvalColumn.dummy.name]
                     ]
                     if leaderboard_df.empty is False
                     else leaderboard_df
                 ),
-                headers=[c.name for c in fields(AutoEvalColumn) if c.never_hidden] + shown_columns.value,
+                headers=[c.name for c in fields(AutoEvalColumn) if c.never_hidden] + shown_columns.value + benchmark_columns,
                 datatype=TYPES,
                 elem_id="leaderboard-table",
                 interactive=False,
@@ -313,7 +326,7 @@ with demo:
             demo.load(load_query, inputs=[], outputs=[search_bar])
 
             for selector in [shown_columns, filter_columns_type, filter_columns_precision, filter_columns_size]:
-                selector.select(
+                selector.change(
                     update_table,
                     [
                         hidden_leaderboard_table_for_search,

backend-cli.py

@@ -17,7 +17,7 @@ from src.backend.manage_requests import EvalRequest
 from src.leaderboard.read_evals import EvalResult
 
 from src.envs import QUEUE_REPO, RESULTS_REPO, API, DEBUG_QUEUE_REPO, DEBUG_RESULTS_REPO
-from src.utils import my_snapshot_download, analyze_gpu_stats, parse_nvidia_smi, monitor_gpus
+from src.utils import my_snapshot_download, analyze_gpu_stats, parse_nvidia_smi, monitor_gpus, get_gpu_details
 
 from src.leaderboard.read_evals import get_raw_eval_results
 
@@ -28,6 +28,8 @@ import time
 import pprint
 import logging
 
+from lm_eval.filters.extraction import RegexFilter
+
 
 # Configure the root logger
 logging.basicConfig(
@@ -42,6 +44,20 @@ eval_logger = logging.getLogger("lm-eval")
 # Explicitly set the level for 'lm-eval' logger to WARNING
 eval_logger.setLevel(logging.WARNING)
 
+def tuple_input_decorator(func):
+    def wrapper(self, resps, docs):
+        stripped_resps = [[resp_data[0] for resp_data in group] for group in resps]
+
+        filtered_resps = func(self, stripped_resps, docs)
+        
+        combined_resps = []
+        for original_group, new_group in zip(resps, filtered_resps):
+            combined_group = [(new_resp,) + rest_of_data[1:] for new_resp, rest_of_data in zip(new_group, original_group)]
+            combined_resps.append(combined_group)
+
+        return combined_resps
+    return wrapper
+
 
 def my_set_eval_request(api, eval_request, set_to_status, hf_repo, local_dir):
     for i in range(10):
@@ -126,9 +142,6 @@ def request_to_result_name(request: EvalRequest) -> str:
 def process_evaluation(task: Task, eval_request: EvalRequest, limit: Optional[int] = None) -> dict:
     batch_size = 1
     batch_size = eval_request.batch_size
-    
-    if args.debug:
-        RESULTS_REPO = DEBUG_RESULTS_REPO
 
     init_gpu_info = analyze_gpu_stats(parse_nvidia_smi())
     # if init_gpu_info['Mem(M)'] > 500:
@@ -137,6 +150,12 @@ def process_evaluation(task: Task, eval_request: EvalRequest, limit: Optional[in
     stop_event = threading.Event()
     monitor_thread = threading.Thread(target=monitor_gpus, args=(stop_event, 5, gpu_stats_list))
     monitor_thread.start()
+    
+    original_apply = RegexFilter.apply
+    if task.benchmark in ["gsm8k", "gsm8k_cot", "gsm8k_cot_self_consistency", "gsm8k_custom"]:
+        RegexFilter.apply = tuple_input_decorator(RegexFilter.apply)
+    else:
+        RegexFilter.apply = original_apply
 
     try:
         results = run_evaluation(
@@ -198,6 +217,8 @@ def process_evaluation(task: Task, eval_request: EvalRequest, limit: Optional[in
         repo_id=RESULTS_REPO,
         repo_type="dataset",
     )
+    
+    RegexFilter.apply = original_apply
     return results
 
 
@@ -366,21 +387,7 @@ def maybe_refresh_results(thr: int, hard_task_lst: Optional[list[str]] = None) -
 
     return False
 
-
-def get_gpu_details():
-    gpus = GPUtil.getGPUs()
-    gpu = gpus[0]
-    name = gpu.name.replace(" ", "-")
-    # Convert memory from MB to GB and round to nearest whole number
-    memory_gb = round(gpu.memoryTotal / 1024)
-    memory = f"{memory_gb}GB"
-    formatted_name = f"{name}-{memory}"
-    return formatted_name
-
 def process_pending_requests() -> bool:
-    if args.debug:
-        QUEUE_REPO = DEBUG_QUEUE_REPO
-        
     sanity_checks()
     print("Processing pending requests")
     current_pending_status = [PENDING_STATUS]
@@ -443,13 +450,14 @@ def get_args():
     parser = argparse.ArgumentParser(description="Run the backend")
     parser.add_argument("--debug", action="store_true", help="Run in debug mode")
     # debug parameters
-    parser.add_argument("--task", type=str, default="selfcheckgpt,mmlu", help="Task to debug")
+    parser.add_argument("--task", type=str, default="selfcheckgpt,mmlu, gsm8k", help="Task to debug")
     parser.add_argument("--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1,mistralai/Mixtral-8x7B-v0.1", help="Model to debug")
     parser.add_argument("--precision", type=str, default="float32,float16,8bit,4bit", help="Precision to debug")
     parser.add_argument("--inference-framework", type=str, default="hf-chat", help="Inference framework to debug")
     parser.add_argument("--limit", type=int, default=None, help="Limit for the number of samples")
     parser.add_argument("--gpu-type", type=str, default="NVIDIA-A100-PCIe-80GB", 
                         help="GPU type. NVIDIA-A100-PCIe-80GB; NVIDIA-RTX-A5000-24GB; NVIDIA-H100-PCIe-80GB")
+    parser.add_argument("--debug_repo", action="store_true", help="Use debug repo")
     return parser.parse_args()
 
 
@@ -457,7 +465,7 @@ if __name__ == "__main__":
     args = get_args()
     local_debug = args.debug
     # debug specific task by ping
-    if local_debug:
+    if local_debug and not args.debug_repo:
         # debug_model_names = [args.model]  # Use model from arguments
         # debug_task_name = [args.task]  # Use task from arguments
         debug_model_names = args.model.split(",")
@@ -471,42 +479,60 @@ if __name__ == "__main__":
                     task_name = task.benchmark
                     if task_name not in debug_task_name:
                         continue
-                    try:
-                        eval_request = EvalRequest(
-                            model=debug_model_name,
-                            private=False,
-                            status="",
-                            json_filepath="",
-                            precision=precision,  # Use precision from arguments
-                            inference_framework=args.inference_framework,  # Use inference framework from arguments
-                            gpu_type=args.gpu_type
-                        )
-                        curr_gpu_type = get_gpu_details()
-                        if eval_request.gpu_type != curr_gpu_type:
-                            print(f"GPU type mismatch: {eval_request.gpu_type} vs {curr_gpu_type}")
-                            raise Exception("GPU type mismatch")
-                        results = process_evaluation(task, eval_request, limit=args.limit)
-                    except Exception as e:
-                        print(f"debug running error: {e}")
-    else:
+                    # try:
+                    eval_request = EvalRequest(
+                        model=debug_model_name,
+                        private=False,
+                        status="",
+                        json_filepath="",
+                        precision=precision,  # Use precision from arguments
+                        inference_framework=args.inference_framework,  # Use inference framework from arguments
+                        gpu_type=args.gpu_type
+                    )
+                    curr_gpu_type = get_gpu_details()
+                    if eval_request.gpu_type != curr_gpu_type:
+                        print(f"GPU type mismatch: {eval_request.gpu_type} vs {curr_gpu_type}")
+                        raise Exception("GPU type mismatch")
+                    results = process_evaluation(task, eval_request, limit=args.limit)
+                    # except Exception as e:
+                    #     print(f"debug running error: {e}")
+    elif local_debug and args.debug_repo:
+        QUEUE_REPO = DEBUG_QUEUE_REPO
+        RESULTS_REPO = DEBUG_RESULTS_REPO
         while True:
             res = False
-
             # if random.randint(0, 10) == 0:
             res = process_pending_requests()
             print(f"waiting for 60 seconds")
             time.sleep(60)
-
             # if res is False:
             #     if random.randint(0, 5) == 0:
             #         res = maybe_refresh_results(100)
             #     else:
             #         res = process_finished_requests(100)
-
             # time.sleep(60)
-
             # if res is False:
             #     if random.randint(0, 5) == 0:
             #         res = maybe_refresh_results(0)
             #     else:
             #         res = process_finished_requests(0)
+    elif not local_debug and not args.debug_repo:
+        while True:
+           res = False
+           # if random.randint(0, 10) == 0:
+           res = process_pending_requests()
+           print(f"waiting for 60 seconds")
+           time.sleep(60)
+           # if res is False:
+           #     if random.randint(0, 5) == 0:
+           #         res = maybe_refresh_results(100)
+           #     else:
+           #         res = process_finished_requests(100)
+           # time.sleep(60)
+           # if res is False:
+           #     if random.randint(0, 5) == 0:
+           #         res = maybe_refresh_results(0)
+           #     else:
+           #         res = process_finished_requests(0)
+    else:
+        raise Exception("Cannot use debug_repo without local debug flag")

requirements.txt

@@ -30,4 +30,5 @@ evaluate
 spacy==3.7.4
 selfcheckgpt
 immutabledict
-gputil
+gputil
+bitsandbytes

src/backend/envs.py

@@ -57,6 +57,7 @@ class Tasks(Enum):
 
     # task20 = Task("race", "acc", "RACE", 0)
     task21 = Task("mmlu", "acc", "MMLU", 5)
+    task22 = Task("gsm8k_custom", "em", "GSM8K", 5)
 
 
 EVAL_REQUESTS_PATH_BACKEND = os.path.join(CACHE_PATH, "eval-queue-bk")

src/backend/hflm_with_measurement.py

@@ -295,6 +295,8 @@ class HFLMWithMeasurement(HFLM):
         # and we don't want a warning from HF
         generation_kwargs["temperature"] = generation_kwargs.get("temperature", 0.0)
         do_sample = generation_kwargs.get("do_sample", None)
+        
+        is_gsm8k = generation_kwargs.get("is_gsm8k", False)
 
         # The temperature has to be a strictly positive float -- if it is 0.0, use greedy decoding strategies
         if generation_kwargs.get("temperature") == 0.0 and do_sample is None:
@@ -302,22 +304,40 @@ class HFLMWithMeasurement(HFLM):
 
         if do_sample is False and generation_kwargs.get("temperature") == 0.0:
             generation_kwargs.pop("temperature")
+        
+        generation_kwargs.pop("is_gsm8k")
+
+        if not is_gsm8k:
         # build stopping criteria
-        stopping_criteria = stop_sequences_criteria(
-            self.tokenizer, stop, context.shape[1], context.shape[0]
-        )
-        stop_watch = StopWatch(self.tokenizer)
-        start = time()
-        res = self.model.generate(
-            input_ids=context,
-            max_length=max_length,
-            stopping_criteria=stopping_criteria,
-            pad_token_id=self.tokenizer.pad_token_id,
-            use_cache=True,
-            streamer=stop_watch,
-            **generation_kwargs,
-        )
-        end = time()
+            stopping_criteria = stop_sequences_criteria(
+                self.tokenizer, stop, context.shape[1], context.shape[0]
+            )
+            stop_watch = StopWatch(self.tokenizer)
+            start = time()
+            res = self.model.generate(
+                input_ids=context,
+                max_length=max_length,
+                stopping_criteria=stopping_criteria,
+                pad_token_id=self.tokenizer.pad_token_id,
+                use_cache=True,
+                streamer=stop_watch,
+                **generation_kwargs,
+            )
+            end = time()
+        else:
+            # print("Using GSM8K")
+            stop_watch = StopWatch(self.tokenizer)
+            start = time()
+            res = self.model.generate(
+                input_ids=context,
+                max_length=max_length,
+                eos_token_id=stop,
+                pad_token_id=self.tokenizer.pad_token_id,
+                use_cache=True,
+                streamer=stop_watch,
+                **generation_kwargs,
+            )
+            end = time()
 
         batch_size = context.shape[0]
         output_length = stop_watch.decoding_iterations
@@ -408,6 +428,11 @@ class HFLMWithMeasurement(HFLM):
                 until = [eos]
             else:
                 until.append(eos)
+            
+            is_gsm8k = kwargs.get("is_gsm8k", False)
+            if is_gsm8k:
+                until = [self.tokenizer.eos_token_id, self.tokenizer.convert_tokens_to_ids("<|eot_id|>")]
+                    
             if "max_gen_toks" in kwargs.keys():
                 max_gen_toks = kwargs.pop("max_gen_toks")
             else:
@@ -427,6 +452,8 @@ class HFLMWithMeasurement(HFLM):
                 left_truncate_len=max_ctx_len,
                 truncation=self.truncation,
             )
+            
+            # print("context: ", self.tok_decode(context_enc[0]))
             context_enc = context_enc.to(self.device)
             attn_masks = attn_masks.to(self.device)
 
@@ -445,16 +472,18 @@ class HFLMWithMeasurement(HFLM):
             for cont_toks, context in zip(cont_toks_list, contexts):
                 # discard context + left-padding toks if using causal decoder-only LM
                 if self.AUTO_MODEL_CLASS == transformers.AutoModelForCausalLM:
+                    # print("After Generation: ", self.tok_decode(cont_toks))
                     cont_toks = cont_toks[context_enc.shape[1] :]
-
+                
                 s = self.tok_decode(cont_toks)
 
                 # use secondary stop seqs to cut off should-have-been-stopped content post-hoc
-                for term in until:
-                    if len(term) > 0:
-                        # ignore '' separator,
-                        # for seq2seq case where self.tok_decode(self.eot_token_id) = ''
-                        s = s.split(term)[0]
+                if not is_gsm8k:
+                    for term in until:
+                        if len(term) > 0:
+                            # ignore '' separator,
+                            # for seq2seq case where self.tok_decode(self.eot_token_id) = ''
+                            s = s.split(term)[0]
 
                 res.append((s, end_to_end_time, prefilling_time, token_per_sec))
 

src/backend/tasks/gsm8k/gsm8k-custom.yaml

@@ -0,0 +1,44 @@
+group:
+  - math_word_problems
+task: gsm8k_custom
+dataset_path: gsm8k
+dataset_name: main
+output_type: generate_until
+training_split: train
+fewshot_split: train
+test_split: test
+doc_to_text: "Question: {{question}}\nAnswer:"
+doc_to_target: "{{answer}}" #" {{answer.split('### ')[-1].rstrip()}}"
+metric_list:
+  - metric: exact_match
+    aggregation: mean
+    higher_is_better: true
+    ignore_case: true
+    ignore_punctuation: false
+    regexes_to_ignore:
+      - ","
+      - "\\$"
+      - "(?s).*#### "
+      - "\\.$"
+generation_kwargs:
+  until:
+    - "<|eot_id|>"
+  do_sample: false
+  temperature: 0.0
+  is_gsm8k: true
+repeats: 1
+num_fewshot: 5
+filter_list:
+  # - name: "strict-match"
+  #   filter:
+  #     - function: "regex"
+  #       regex_pattern: "#### (\\-?[0-9\\.\\,]+)"
+  #     - function: "take_first"
+  - name: "flexible-extract"
+    filter:
+      - function: "regex"
+        group_select: -1
+        regex_pattern: "(-?[$0-9.,]{2,})|(-?[0-9]+)"
+      - function: "take_first"
+metadata:
+  version: 3.0

src/display/utils.py

@@ -75,6 +75,7 @@ class Tasks(Enum):
     # # XXX include me back at some point
     selfcheck = Task("selfcheckgpt", "max-selfcheckgpt", "SelfCheckGPT")
     mmlu = Task("mmlu", "acc", "MMLU") #MMLU/Acc (5-shot)
+    gsm8k = Task("gsm8k_custom", "em", "GSM8K") #GSM8K/EM (8-shot)
 
 
 # These classes are for user facing column names,
@@ -104,16 +105,16 @@ auto_eval_column_dict.append(["inference_framework", ColumnContent, ColumnConten
 for task in Tasks:
     auto_eval_column_dict.append([task.name, ColumnContent, ColumnContent(task.value.col_name, "number", True)])
     # System performance metrics
-    auto_eval_column_dict.append([f"{task.name}_end_to_end_time", ColumnContent, ColumnContent(f"{task.value.col_name} {E2Es}", "number", True)])
-    auto_eval_column_dict.append([f"{task.name}_batch_size", ColumnContent, ColumnContent(f"{task.value.col_name} {BATCH_SIZE}", "number", True)])
-    # auto_eval_column_dict.append([f"{task.name}_precision", ColumnContent, ColumnContent(f"{task.value.col_name} {PRECISION}", "str", True)])
-    auto_eval_column_dict.append([f"{task.name}_gpu_mem", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Mem}", "number", True)])
-    auto_eval_column_dict.append([f"{task.name}_gpu", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Name}", "str", True)])
-    auto_eval_column_dict.append([f"{task.name}_gpu_util", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Util}", "number", True)])
+    auto_eval_column_dict.append([f"{task.name}_end_to_end_time", ColumnContent, ColumnContent(f"{task.value.col_name} {E2Es}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_batch_size", ColumnContent, ColumnContent(f"{task.value.col_name} {BATCH_SIZE}", "number", True, hidden=True)])
+    # auto_eval_column_dict.append([f"{task.name}_precision", ColumnContent, ColumnContent(f"{task.value.col_name} {PRECISION}", "str", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_gpu_mem", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Mem}", "number", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_gpu", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Name}", "str", True, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_gpu_util", ColumnContent, ColumnContent(f"{task.value.col_name} {GPU_Util}", "number", True, hidden=True)])
     if task.value.benchmark in MULTIPLE_CHOICEs:
         continue
-    # auto_eval_column_dict.append([f"{task.name}_prefilling_time", ColumnContent, ColumnContent(f"{task.value.col_name} {PREs}", "number", False)])
-    auto_eval_column_dict.append([f"{task.name}_decoding_throughput", ColumnContent, ColumnContent(f"{task.value.col_name} {TS}", "number", True)])
+    # auto_eval_column_dict.append([f"{task.name}_prefilling_time", ColumnContent, ColumnContent(f"{task.value.col_name} {PREs}", "number", False, hidden=True)])
+    auto_eval_column_dict.append([f"{task.name}_decoding_throughput", ColumnContent, ColumnContent(f"{task.value.col_name} {TS}", "number", True, hidden=True)])
 
 
 # Model information
@@ -242,8 +243,8 @@ class Precision(Enum):
 
 
 # Column selection
-COLS = [c.name for c in fields(AutoEvalColumn) if not c.hidden]
-TYPES = [c.type for c in fields(AutoEvalColumn) if not c.hidden]
+COLS = [c.name for c in fields(AutoEvalColumn)]
+TYPES = [c.type for c in fields(AutoEvalColumn)]
 COLS_LITE = [c.name for c in fields(AutoEvalColumn) if c.displayed_by_default and not c.hidden]
 TYPES_LITE = [c.type for c in fields(AutoEvalColumn) if c.displayed_by_default and not c.hidden]
 

src/submission/check_validity.py

@@ -130,7 +130,8 @@ def already_submitted_models(requested_models_dir: str) -> set[str]:
                     continue
                 with open(os.path.join(root, file), "r") as f:
                     info = json.load(f)
-                    file_names.append(f"{info['model']}_{info['revision']}_{info['precision']}_{info['inference_framework']}_{info['gpu_type']}")
+                    if not info["status"] == "FINISHED" and not info["status"] == "RUNNING":
+                        file_names.append(f"{info['model']}_{info['revision']}_{info['precision']}_{info['inference_framework']}_{info['gpu_type']}")
 
                     # Select organisation
                     if info["model"].count("/") == 0 or "submitted_time" not in info:

src/utils.py

@@ -3,12 +3,48 @@ from huggingface_hub import snapshot_download
 import subprocess
 import re
 import os
+import GPUtil
 
 try:
     from src.display.utils import GPU_TEMP, GPU_Mem, GPU_Power, GPU_Util, GPU_Name
 except:
     print("local debug: from display.utils")
     from display.utils import GPU_TEMP, GPU_Mem, GPU_Power, GPU_Util, GPU_Name
+    
+MEM_BW_DICT ={
+    "NVIDIA-A100-PCIe-80GB": 1935,
+    "NVIDIA-A100-SXM-80GB": 2039,
+    "NVIDIA-H100-PCIe-80GB": 2039,
+    "NVIDIA-RTX-A5000-24GB": 768
+}
+
+PEAK_FLOPS_DICT = {
+    "float32":{
+        "NVIDIA-A100-PCIe-80GB": 312e12,
+        "NVIDIA-A100-SXM-80GB": 312e12,
+        "NVIDIA-H100-PCIe-80GB": 756e12,
+        "NVIDIA-RTX-A5000-24GB": 222.2e12
+    },
+    "float16":{
+        "NVIDIA-A100-PCIe-80GB": 624e12,
+        "NVIDIA-A100-SXM-80GB": 624e12,
+        "NVIDIA-H100-PCIe-80GB": 1513e12,
+        "NVIDIA-RTX-A5000-24GB": 444.4e12
+    },
+    "8bit":{
+        "NVIDIA-A100-PCIe-80GB": 1248e12,
+        "NVIDIA-A100-SXM-80GB": 1248e12,
+        "NVIDIA-H100-PCIe-80GB": 3026e12,
+        "NVIDIA-RTX-A5000-24GB": 889e12
+    },
+    "4bit": {
+        "NVIDIA-A100-PCIe-80GB": 2496e12,
+        "NVIDIA-A100-SXM-80GB": 2496e12,
+        "NVIDIA-H100-PCIe-80GB": 6052e12,
+        "NVIDIA-RTX-A5000-24GB": 1778e12
+    }
+
+}
 
 def my_snapshot_download(repo_id, revision, local_dir, repo_type, max_workers):
     for i in range(10):
@@ -52,11 +88,11 @@ def parse_nvidia_smi():
             print("Failed to query GPU indices.")
             return []
         gpu_indices = result.stdout.strip().split('\n')
-    print(f"gpu_indices: {gpu_indices}")
+    # print(f"gpu_indices: {gpu_indices}")
     gpu_stats = []
 
     gpu_info_pattern = re.compile(r'(\d+)C\s+P\d+\s+(\d+)W / \d+W\s+\|\s+(\d+)MiB / \d+MiB\s+\|\s+(\d+)%')
-    gpu_name_pattern = re.compile(r'NVIDIA\s+([\w\s]+?\d+GB)')
+    gpu_name_pattern = re.compile(r'NVIDIA\s+([\w\s]+\d+(?:\s*GB)?)')
 
     gpu_name = ""
     for index in gpu_indices:
@@ -80,7 +116,7 @@ def parse_nvidia_smi():
 
             if len(gpu_info) >= 4:
                 gpu_stats.append(gpu_info)
-    print(f"gpu_stats: {gpu_stats}")
+    # print(f"gpu_stats: {gpu_stats}")
     gpu_name = f"{len(gpu_stats)}x{gpu_name}"
     gpu_stats_total = {
                         GPU_TEMP: 0,
@@ -131,5 +167,70 @@ def analyze_gpu_stats(stats_list):
 
     return avg_stats
 
+def get_gpu_number():
+    visible_devices = os.getenv('CUDA_VISIBLE_DEVICES', None)
+    if visible_devices is not None:
+        gpu_indices = visible_devices.split(',')
+    else:
+        # Query all GPU indices if CUDA_VISIBLE_DEVICES is not set
+        result = subprocess.run(['nvidia-smi', '--query-gpu=index', '--format=csv,noheader'], capture_output=True, text=True)
+        if result.returncode != 0:
+            print("Failed to query GPU indices.")
+            return []
+        gpu_indices = result.stdout.strip().split('\n')
+    # print(f"gpu_indices: {gpu_indices}")
+    gpu_stats = []
+
+    gpu_info_pattern = re.compile(r'(\d+)C\s+P\d+\s+(\d+)W / \d+W\s+\|\s+(\d+)MiB / \d+MiB\s+\|\s+(\d+)%')
+
+    for index in gpu_indices:
+        result = subprocess.run(['nvidia-smi', '-i', index], capture_output=True, text=True)
+        output = result.stdout.strip()
+        lines = output.split("\n")
+        for line in lines:
+            match = gpu_info_pattern.search(line)
+            gpu_info = {}
+            if match:
+                temp, power_usage, mem_usage, gpu_util = map(int, match.groups())
+                gpu_info.update({
+                    GPU_TEMP: temp,
+                    GPU_Power: power_usage,
+                    GPU_Mem: round(mem_usage / 1024, 2),
+                    GPU_Util: gpu_util
+                })
+
+            if len(gpu_info) >= 4:
+                gpu_stats.append(gpu_info)
+                
+    return len(gpu_stats)
+
+def get_gpu_details():
+    gpus = GPUtil.getGPUs()
+    gpu = gpus[0]
+    name = gpu.name.replace(" ", "-")
+    # Convert memory from MB to GB and round to nearest whole number
+    memory_gb = round(gpu.memoryTotal / 1024)
+    memory = f"{memory_gb}GB"
+    formatted_name = f"{name}-{memory}"
+    return formatted_name
+
+def get_peak_bw(gpu_name):
+    return MEM_BW_DICT[gpu_name]
+
+def get_peak_flops(gpu_name, precision):
+    return PEAK_FLOPS_DICT[precision][gpu_name]
+
+def transfer_precision2bytes(precision):
+    if precision == "float32":
+        return 4
+    elif precision == "float16":
+        return 2
+    elif precision == "8bit":
+        return 1
+    elif precision == "4bit":
+        return 0.5
+    else:
+        raise ValueError(f"Unsupported precision: {precision}")
+
 if __name__ == "__main__":
     print(analyze_gpu_stats(parse_nvidia_smi()))