this can run on local and you may need to specify if 'model' not in args.__dict__ or len(args.model)<2:

args.model='/home/sribd/chenghao/models/phi-2'
# args.model='microsoft/phi-2'
here in order to run online

.gitignore

@@ -0,0 +1 @@
+get_loss/__pycache__/

app.py

@@ -15,13 +15,18 @@ def color_text(text_list=["hi", "FreshEval","!"], loss_list=[0.1,0.7]):
     根据损失值为文本着色。
     """
     highlighted_text = []
+    # print('loss_list',loss_list)
+    # ndarray to list
+    loss_list = loss_list.tolist()
     loss_list=[0]+loss_list
+    # print('loss_list',loss_list)
+    # print('text_list',text_list)
+    pdb.set_trace()
     for text, loss in zip(text_list, loss_list):
         # color = "#FF0000" if float(loss) > 0.5 else "#00FF00"
-        color=loss/25
+        color=loss/20#TODO rescale
         # highlighted_text.append({"text": text, "bg_color": color})
         highlighted_text.append((text, color))
-
         print('highlighted_text',highlighted_text)
     return highlighted_text
 
@@ -32,7 +37,7 @@ def get_text(ids_list=[0.1,0.7], tokenizer=None):
     """
     # return ['Hi', 'Adam']
     # tokenizer = AutoTokenizer.from_pretrained(tokenizer)
-    print('ids_list',ids_list)
+    # print('ids_list',ids_list)
     # pdb.set_trace()
     text=[]
     for id in ids_list:
@@ -64,6 +69,7 @@ def color_pipeline(texts=["Hi","FreshEval","!"],  model=None):
     # pdb.set_trace()
     # {'logit':logit,'input_ids':input_chunk,'tokenizer':tokenizer,'neg_log_prob_temp':neg_log_prob_temp}
     ids, loss =rtn_dic['input_ids'],rtn_dic['loss']#= get_ids_loss(text, tokenizer, model)
+    # notice here is numpy ndarray
     tokenizer=rtn_dic['tokenizer'] # get tokenizer
     text = get_text(ids, tokenizer)
     # print('ids, loss ,text',ids, loss ,text)

get_loss/get_loss_hf.py

@@ -123,7 +123,6 @@ def print_model_parameters_in_billions(model):
 
 
 def load_hf_model(path, cache_path):
-    hf_tokenizer = AutoTokenizer.from_pretrained(path)
     if cache_path is not None:
         # pdb.set_trace()
         hf_model = AutoModelForCausalLM.from_pretrained(path,
@@ -134,6 +133,7 @@ def load_hf_model(path, cache_path):
         hf_model = AutoModelForCausalLM.from_pretrained(path,
                                                         device_map=device,
                                                         trust_remote_code=True).eval()
+    hf_tokenizer = AutoTokenizer.from_pretrained(path)
 
     print_model_parameters_in_billions(hf_model)
 
@@ -253,11 +253,12 @@ def run_get_loss(args=None):
     if 'model_type' not in args.__dict__:
         args.model_type='hf'
     if 'model' not in args.__dict__ or len(args.model)<2:
-        # args.model='/home/sribd/chenghao/models/phi-2'
-        args.model='microsoft/phi-2'
+        args.model='/home/sribd/chenghao/models/phi-2'
+        # args.model='microsoft/phi-2'
 
     if 'model_cache' not in args.__dict__:
-        args.model_cache=args.model
+        # args.model_cache=args.model
+        args.model_cache=None
 
     # args = parser.parse_args()