From 112108d564e4834692796ad82dd5ea9153ce47fd Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Sat, 18 Nov 2023 19:39:23 +0800
Subject: [PATCH] fix evaluator and cached_file in 4.31.0

Former-commit-id: ff6056405dea8e89a95fd3741fd309d3c7679896
---
 src/llmtuner/eval/evaluator.py | 14 +++++++++++---
 src/llmtuner/model/utils.py    | 35 +++++++++++++++++++++++++---------
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/llmtuner/eval/evaluator.py b/src/llmtuner/eval/evaluator.py
index b2e04bec..1fbd40ee 100644
--- a/src/llmtuner/eval/evaluator.py
+++ b/src/llmtuner/eval/evaluator.py
@@ -3,6 +3,7 @@
 import os
 import json
 import torch
+import inspect
 import tiktoken
 import numpy as np
 from tqdm import tqdm, trange
@@ -45,13 +46,18 @@ class Evaluator:
         return [chr(ord("A") + offset.item()) for offset in torch.argmax(choice_probs, dim=-1)]
 
     def eval(self) -> None:
+        if "token" in inspect.signature(cached_file).parameters:
+            kwargs = {"token": self.model_args.hf_hub_token}
+        elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
+            kwargs = {"use_auth_token": self.model_args.hf_hub_token}
+
         mapping = cached_file(
             path_or_repo_id = os.path.join(self.eval_args.task_dir, self.eval_args.task),
             filename="mapping.json",
             cache_dir=self.model_args.cache_dir,
-            token=self.model_args.hf_hub_token,
-            revision=self.model_args.model_revision
+            **kwargs
         )
+
         with open(mapping, "r", encoding="utf-8") as f:
             categorys: Dict[str, Dict[str, str]] = json.load(f)
 
@@ -62,7 +68,9 @@ class Evaluator:
             dataset = load_dataset(
                 path=os.path.join(self.eval_args.task_dir, self.eval_args.task),
                 name=subject,
-                download_mode="force_redownload"
+                cache_dir=self.model_args.cache_dir,
+                download_mode=self.eval_args.download_mode,
+                token=self.model_args.hf_hub_token
             )
             pbar.set_postfix_str(categorys[subject]["name"])
             inputs, outputs, labels = [], [], []
diff --git a/src/llmtuner/model/utils.py b/src/llmtuner/model/utils.py
index e7445f8d..5c0b1655 100644
--- a/src/llmtuner/model/utils.py
+++ b/src/llmtuner/model/utils.py
@@ -1,4 +1,5 @@
 import torch
+import inspect
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
 
 from transformers.utils import cached_file
@@ -94,19 +95,35 @@ def load_valuehead_params(
     """
     kwargs = {
         "path_or_repo_id": path_or_repo_id,
-        "cache_dir": model_args.cache_dir,
-        "token": model_args.hf_hub_token
+        "cache_dir": model_args.cache_dir
     }
+
+    if "token" in inspect.signature(cached_file).parameters:
+        kwargs["token"] = model_args.hf_hub_token
+    elif "use_auth_token" in inspect.signature(cached_file).parameters: # for transformers==4.31.0
+        kwargs["use_auth_token"] = model_args.hf_hub_token
+    else:
+        logger.warning("Ignore `hf_hub_token` since matched parameter is not found.")
+
     try:
         vhead_file = cached_file(filename=WEIGHTS_NAME, **kwargs)
-    except:
-        try:
-            vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
-        except:
-            logger.warning("Provided path ({}) does not contain valuehead weights.".format(path_or_repo_id))
-            return None
+        return torch.load(vhead_file, map_location="cpu")
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(WEIGHTS_NAME, str(err)))
 
-    return torch.load(vhead_file, map_location="cpu")
+    try:
+        from safetensors import safe_open
+        vhead_file = cached_file(filename=SAFE_WEIGHTS_NAME, **kwargs)
+        with safe_open(vhead_file, framework="pt", device="cpu") as f:
+            return {
+                "v_head.summary.weight": f.get_tensor("v_head.summary.weight"),
+                "v_head.summary.bias": f.get_tensor("v_head.summary.bias")
+            }
+    except Exception as err:
+        logger.info("Failed to load {}: {}".format(SAFE_WEIGHTS_NAME, str(err)))
+
+    logger.warning("Provided path ({}) does not contain valuehead weights.".format(path_or_repo_id))
+    return None
 
 
 def prepare_model_for_training(