move to sribd

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
 import os
-# from transformers import AutoTokenizer
-os.system('git clone https://github.com/EleutherAI/lm-evaluation-harness')
-os.system('cd lm-evaluation-harness')
-os.system('pip install -e .')
+from transformers import AutoTokenizer
+from .get_loss.get_loss_hf import run_get_loss
+
+# os.system('git clone https://github.com/EleutherAI/lm-evaluation-harness')
+# os.system('cd lm-evaluation-harness')
+# os.system('pip install -e .')
 # 第一个功能：基于输入文本和对应的损失值对文本进行着色展示
+
 def color_text(text_list=["hi", "FreshEval"], loss_list=[0.1,0.7]):
     """
     根据损失值为文本着色。
@@ -45,6 +48,8 @@ def color_pipeline(text=["hi", "FreshEval"],  model=None):
     """
     给定一个文本，返回其对应的着色文本。
     """
+    # rtn_dic=run_get_loss()
+    # {'logit':logit,'input_ids':input_chunk,'tokenizer':tokenizer,'neg_log_prob_temp':neg_log_prob_temp}
     tokenizer=None # get tokenizer
     ids, loss = get_ids_loss(text, tokenizer, model)
     text = get_text(ids, tokenizer)
@@ -61,7 +66,7 @@ with gr.Blocks() as demo:
             # TODO craw and drop the file
 
             # loss_input = gr.Number(label="loss")
-            model_input = gr.Textbox(label="model name", placeholder="input your model name here...")
+            model_input = gr.Textbox(label="model name", placeholder="input your model name here... now I am trying phi-2...")
             # TODO select models that can be used online
             # TODO maybe add our own models
 

get_loss/get_loss.py

@@ -9,6 +9,8 @@ import gc
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime
 import argparse
+import mamba_ssm
+import rwkv
 
 
 RWKV4_TOKENIZER_FILE = "./support/20B_tokenizer.json"
@@ -142,13 +144,13 @@ def eval_rwkv(model, tokenizer, texts, chunk_size, v4pile=False):
             input_length = len(input_seq)
 
             neg_log_prob_temp = 0
-            for begin in range(0, input_length, chunk_size):
-                input_chunk = input_seq[begin: begin + chunk_size]
+            # for begin in range(0, input_length, chunk_size):
+            input_chunk = input_seq[:chunk_size]
 
-                logit = model.forward(input_chunk, None, full_output=True)[0]
+            logit = model.forward(input_chunk, None, full_output=True)[0]
 
-                if len(input_chunk) == 1:
-                    logit = logit.unsqueeze(0)
+            if len(input_chunk) == 1:
+                logit = logit.unsqueeze(0)
 
     #             log_sum = calculate_log_sum(logit, torch.tensor(input_chunk).cuda())
 
@@ -165,7 +167,7 @@ def eval_rwkv(model, tokenizer, texts, chunk_size, v4pile=False):
     # print(f'log probability sum: {sum(rwkv_test_data) / len(rwkv_test_data):.2f}')
     # print(f'avg tokens: {sum(rwkv_token_length_list) / len(rwkv_token_length_list):.0f}')
 
-    return logit
+    return logit,logit,input_chunk,tokenizer
 
 
 def eval_hf_model(model, tokenizer, texts, chunk_size):
@@ -183,9 +185,9 @@ def eval_hf_model(model, tokenizer, texts, chunk_size):
 
             neg_log_prob_temp = 0
             # for begin in range(0, seq_length, chunk_size):
-                input_chunk = inputs['input_ids'][:, begin: begin + chunk_size]
+            input_chunk = inputs['input_ids'][:, :chunk_size]
 
-                logit = model.forward(input_ids=input_chunk).logits[0, :, :]
+            logit = model.forward(input_ids=input_chunk).logits[0, :, :]
 
     #             log_sum = calculate_log_sum(logit, input_chunk.squeeze(0))
     #             neg_log_prob_temp += log_sum
@@ -201,7 +203,7 @@ def eval_hf_model(model, tokenizer, texts, chunk_size):
     # print(f'log probability sum: {sum(data) / len(data):.2f}')
     # print(f'avg tokens: {sum(token_length_list) / len(token_length_list):.0f}')
 
-    return logit
+    return logit,input_chunk,tokenizer
 
 
 # if __name__ == '__main__':
@@ -236,11 +238,11 @@ def run_get_loss(args):
 
     # eval
     if args.model_type in ['hf', 'mamba']:
-        results = eval_hf_model(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+        return eval_hf_model(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
     elif args.model_type == 'rwkv':
-        results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+        return eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
     elif args.model_type == 'rwkv4pile':
-        results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size, v4pile=True)
+        return eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size, v4pile=True)
     else:
         raise NotImplementedError
 
@@ -252,43 +254,44 @@ def run_get_loss(args):
 
     # print(json.dumps(results, indent=4, ensure_ascii=False))
 
-
+from types import SimpleNamespace
 
 if __name__ == '__main__':
-
-
-
-def run_get_loss(input_string, model_type):
-    # load data
-    texts = [input_string]
-    print(f'data size: {len(texts)}')
-
-    # load model
-    if model_type == 'hf':
-        model, tokenizer = load_hf_model(args.model, args.model_cache)# tokenzier path, model path
-    elif model_type == 'rwkv':
-        model, tokenizer = load_rwkv(args.model)
-    elif model_type == 'mamba':
-        model, tokenizer = load_mamba(args.model)
-    elif model_type == 'rwkv4pile':
-        model, tokenizer = load_rwkv4pile(args.model)
-    else:
-        raise NotImplementedError
-
-    # eval
-    if model_type in ['hf', 'mamba']:
-        results = eval_hf_model(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
-    elif model_type == 'rwkv':
-        results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
-    elif model_type == 'rwkv4pile':
-        results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size, v4pile=True)
-    else:
-        raise NotImplementedError
-
-    results['model_name_or_path'] = args.model
-    results['data_path'] = args.data
-    results['chunk_size'] = args.chunk_size
-
-    make_log(results, args.log_path)
-
-    print(json.dumps(results, indent=4, ensure_ascii=False))
+    args=SimpleNamespace(model='microsft/phi-2',model_type='hf',data='data.json',log_path='./logs/',model_cache=None,chunk_size=1024)
+
+
+
+# def run_get_loss(input_string, model_type):
+#     # load data
+#     texts = [input_string]
+#     print(f'data size: {len(texts)}')
+
+#     # load model
+#     if model_type == 'hf':
+#         model, tokenizer = load_hf_model(args.model, args.model_cache)# tokenzier path, model path
+#     elif model_type == 'rwkv':
+#         model, tokenizer = load_rwkv(args.model)
+#     elif model_type == 'mamba':
+#         model, tokenizer = load_mamba(args.model)
+#     elif model_type == 'rwkv4pile':
+#         model, tokenizer = load_rwkv4pile(args.model)
+#     else:
+#         raise NotImplementedError
+
+#     # eval
+#     if model_type in ['hf', 'mamba']:
+#         results = eval_hf_model(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+#     elif model_type == 'rwkv':
+#         results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+#     elif model_type == 'rwkv4pile':
+#         results = eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size, v4pile=True)
+#     else:
+#         raise NotImplementedError
+
+#     results['model_name_or_path'] = args.model
+#     results['data_path'] = args.data
+#     results['chunk_size'] = args.chunk_size
+
+#     make_log(results, args.log_path)
+
+#     print(json.dumps(results, indent=4, ensure_ascii=False))

get_loss/get_loss_hf.py

@@ -0,0 +1,288 @@
+# import packages
+import os
+# from tqdm import tqdm
+# import warnings
+import json
+import torch.nn.functional as F
+import torch
+import gc
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datetime import datetime
+import argparse
+from types import SimpleNamespace
+
+# import mamba_ssm
+# import rwkv
+
+
+# RWKV4_TOKENIZER_FILE = "./support/20B_tokenizer.json"
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+def load_list_from_json(file_path):
+    """
+    Loads a list of strings from a JSON file.
+
+    :param file_path: Path of the JSON file to be loaded.
+    :return: List of strings loaded from the JSON file.
+    """
+    with open(file_path, 'r', encoding='utf-8') as file:
+        return json.load(file)
+
+
+def calculate_loss(logits, target_token_ids):
+    # shifted_logits = logits[:-1, :]
+    # shifted_targets = target_token_ids[1:]
+
+    # log_probs = F.log_softmax(shifted_logits, dim=-1)
+    loss = torch.nn.functional.cross_entropy(logits[:-1, :].view(-1, logits.shape[-1]), 
+                                             target_token_ids[1:].view(-1), reduction='none')
+
+
+    # target_log_probs = -log_probs.gather(1, shifted_targets.unsqueeze(1)).squeeze()
+    # # print(target_log_probs)
+
+    # log_sum = torch.sum(target_log_probs, dim=-1)
+    # print(perplexity_sum)
+
+    return loss.item()
+
+
+def calculate_log_sum(logits, target_token_ids):
+    shifted_logits = logits[:-1, :]
+    shifted_targets = target_token_ids[1:]
+
+    log_probs = F.log_softmax(shifted_logits, dim=-1)
+
+    target_log_probs = -log_probs.gather(1, shifted_targets.unsqueeze(1)).squeeze()
+    # print(target_log_probs)
+
+    log_sum = torch.sum(target_log_probs, dim=-1)
+    # print(perplexity_sum)
+
+    return log_sum.item()
+
+
+def print_model_parameters_in_billions(model):
+    total_params = sum(p.numel() for p in model.parameters())
+
+    total_params_billion = total_params / 1e9
+
+    print(f"Model parameters: {total_params_billion:.3f} billion")
+
+
+# def make_log(data_dict, folder_path):
+#     if not os.path.exists(folder_path):
+#         try:
+#             os.makedirs(folder_path)
+#             print(f"Directory created at {folder_path}")
+#         except Exception as e:
+#             print(f"Error creating directory: {e}")
+#             return
+
+#     timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+#     file_name = f"{timestamp}.json"
+#     file_path = os.path.join(folder_path, file_name)
+
+#     try:
+#         with open(file_path, 'w') as file:
+#             json.dump(data_dict, file, indent=4)
+#         print(f"Dictionary saved successfully to {file_path}")
+#     except Exception as e:
+#         print(f"Error saving dictionary: {e}")
+
+
+# def load_rwkv(path):
+#     os.environ['RWKV_JIT_ON'] = '1'
+#     os.environ["RWKV_CUDA_ON"] = '1'
+
+#     from rwkv.model import RWKV
+#     from rwkv.utils import PIPELINE
+
+#     rwkv_model = RWKV(model=path, strategy='cuda fp16')
+#     rwkv_pipeline = PIPELINE(rwkv_model, r"rwkv_vocab_v20230424")
+#     rwkv_tokenizer = rwkv_pipeline.tokenizer
+
+#     return rwkv_model, rwkv_tokenizer
+
+
+# def load_rwkv4pile(path):
+#     os.environ['RWKV_JIT_ON'] = '1'
+#     os.environ["RWKV_CUDA_ON"] = '1'
+
+#     from rwkv.model import RWKV
+#     from rwkv.utils import PIPELINE
+
+#     rwkv_model = RWKV(model=path, strategy='cuda fp16')
+#     rwkv_pipeline = PIPELINE(rwkv_model, RWKV4_TOKENIZER_FILE)
+#     rwkv_tokenizer = rwkv_pipeline.tokenizer
+
+#     return rwkv_model, rwkv_tokenizer
+
+
+def load_hf_model(path, cache_path):
+    hf_tokenizer = AutoTokenizer.from_pretrained(path)
+    if cache_path is not None:
+        hf_model = AutoModelForCausalLM.from_pretrained(path,
+                                                        device_map=device,
+                                                        trust_remote_code=True,
+                                                        cache_dir=cache_path).eval()
+    else:
+        hf_model = AutoModelForCausalLM.from_pretrained(path,
+                                                        device_map=device,
+                                                        trust_remote_code=True).eval()
+
+    print_model_parameters_in_billions(hf_model)
+
+    return hf_model, hf_tokenizer
+
+
+# def load_mamba(path):
+#     from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
+
+#     mamba_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+#     mamba_model = MambaLMHeadModel.from_pretrained(path, device="cuda", dtype=torch.float16)
+#     mamba_model.device = torch.device('cuda')
+
+#     print_model_parameters_in_billions(mamba_model)
+
+#     return mamba_model, mamba_tokenizer
+
+
+# def eval_rwkv(model, tokenizer, texts, chunk_size, v4pile=False):
+#     rwkv_test_data = []
+#     rwkv_token_length_list = []
+
+#     for idx, sample in tqdm(enumerate(texts), total=len(texts)):
+
+#         with torch.no_grad():
+
+#             if v4pile:
+#                 input_seq = tokenizer.encode(sample).ids  # v4
+#             else:
+#                 input_seq = tokenizer.encode(sample)
+
+#             input_length = len(input_seq)
+
+#             neg_log_prob_temp = 0
+#             # for begin in range(0, input_length, chunk_size):
+#             input_chunk = input_seq[:chunk_size]
+
+#             logit = model.forward(input_chunk, None, full_output=True)[0]
+
+#             if len(input_chunk) == 1:
+#                 logit = logit.unsqueeze(0)
+
+    #             log_sum = calculate_log_sum(logit, torch.tensor(input_chunk).cuda())
+
+    #             neg_log_prob_temp += log_sum
+
+    #         rwkv_token_length_list.append(input_length)
+    #         rwkv_test_data.append(neg_log_prob_temp)
+
+    # data_dict = {
+    #     'neg_log_prob_sum': sum(rwkv_test_data) / len(rwkv_test_data),
+    #     'avg tokens': sum(rwkv_token_length_list) / len(rwkv_token_length_list),
+    # }
+
+    # print(f'log probability sum: {sum(rwkv_test_data) / len(rwkv_test_data):.2f}')
+    # print(f'avg tokens: {sum(rwkv_token_length_list) / len(rwkv_token_length_list):.0f}')
+
+    return logit,logit,input_chunk,tokenizer
+
+
+def eval_hf_model(model, tokenizer, texts, chunk_size):
+    data = []
+    token_length_list = []
+
+    # for idx, sample in tqdm(enumerate(texts), total=len(texts)):#TODO deleta the forloop
+    with torch.no_grad():
+
+        inputs = tokenizer(texts, return_tensors='pt')
+        inputs = inputs.to(model.device)
+
+        seq_length = inputs['input_ids'].shape[-1]
+
+        neg_log_prob_temp = 0
+        # for begin in range(0, seq_length, chunk_size):
+        input_chunk = inputs['input_ids'][:, :chunk_size]
+
+        logit = model.forward(input_ids=input_chunk).logits[0, :, :]
+
+        log_sum = calculate_log_sum(logit, input_chunk.squeeze(0))# suppose shape of logit is (seq_length, vocab_size),shape of input_chunk is (,seq_length)
+        neg_log_prob_temp += log_sum
+
+        loss = calculate_loss(logit, input_chunk.squeeze(0))
+        neg_log_prob_temp += log_sum
+
+    #         token_length_list.append(seq_length)
+    #         data.append(neg_log_prob_temp)
+
+    # data_dict = {
+    #     'neg_log_prob_sum': sum(data) / len(data),
+    #     'avg tokens': sum(token_length_list) / len(token_length_list),
+    # }
+
+    # print(f'log probability sum: {sum(data) / len(data):.2f}')
+    # print(f'avg tokens: {sum(token_length_list) / len(token_length_list):.0f}')
+    rtn_dic={'logit':logit,'input_ids':input_chunk,'loss':loss,'tokenizer':tokenizer,'neg_log_prob_temp':neg_log_prob_temp}
+    return rtn_dic
+
+
+# if __name__ == '__main__':
+#     parser = argparse.ArgumentParser()
+
+#     parser.add_argument('--model', type=str, required=True, help='model name or path')
+#     parser.add_argument('--model_type', choices=['hf', 'rwkv', 'mamba', 'rwkv4pile'], required=True, help='model type')
+#     parser.add_argument('--data', type=str, required=True, help='data path (json file)')
+#     parser.add_argument('--log_path', type=str, default='./logs/', help='log file path')
+#     parser.add_argument('--model_cache', type=str, help='hugging face model cache')
+#     parser.add_argument('--chunk_size', type=int, default=1024, help='chunk size')
+
+
+def run_get_loss(args):
+    if args is None:
+        args=SimpleNamespace(model='microsoft/phi-2',texts='Hello FreshBench !',model_type='hf',model_cache=None,chunk_size=1024)
+
+    # args = parser.parse_args()
+
+    # load data
+    # texts = load_list_from_json(args.data)
+    texts=args.texts
+    print(f'data size: {len(texts)}')
+
+    # load model
+    if args.model_type == 'hf':
+        model, tokenizer = load_hf_model(args.model, args.model_cache)# tokenzier path, model path
+    # elif args.model_type == 'rwkv':
+    #     model, tokenizer = load_rwkv(args.model)
+    # elif args.model_type == 'mamba':
+    #     model, tokenizer = load_mamba(args.model)
+    # elif args.model_type == 'rwkv4pile':
+    #     model, tokenizer = load_rwkv4pile(args.model)
+    else:
+        raise NotImplementedError
+
+    # eval
+    if args.model_type in ['hf', 'mamba']:
+        return eval_hf_model(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+    # elif args.model_type == 'rwkv':
+    #     return eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size)
+    # elif args.model_type == 'rwkv4pile':
+    #     return eval_rwkv(model=model, tokenizer=tokenizer, texts=texts, chunk_size=args.chunk_size, v4pile=True)
+    else:
+        raise NotImplementedError
+
+    # results['model_name_or_path'] = args.model
+    # results['data_path'] = args.data
+    # results['chunk_size'] = args.chunk_size
+
+    # make_log(results, args.log_path)
+
+    # print(json.dumps(results, indent=4, ensure_ascii=False))
+
+
+if __name__ == '__main__':
+    args=SimpleNamespace(model='microsoft/phi-2',texts='Hello FreshBench !',model_type='hf',model_cache=None,chunk_size=1024)
+    run_get_loss(args)
+
+# run_get_loss(args)