Upload inference_mathglm.py

inference_mathglm.py

@@ -0,0 +1,125 @@
+# -*- encoding: utf-8 -*-
+'''
+@File    :   inference_cogview.py
+@Time    :   2021/10/09 19:41:58
+@Author  :   Ming Ding 
+@Contact :   dm18@mails.tsinghua.edu.cn
+'''
+
+# here put the import lib
+import os
+import sys
+import math
+import random
+import torch
+import argparse
+import stat
+
+from SwissArmyTransformer import mpu, get_args, get_tokenizer
+from SwissArmyTransformer.model import CachedAutoregressiveModel
+from SwissArmyTransformer.generation.sampling_strategies import BaseStrategy
+from SwissArmyTransformer.generation.autoregressive_sampling import filling_sequence
+from SwissArmyTransformer.generation.utils import timed_name, generate_continually
+from SwissArmyTransformer.training import set_random_seed
+
+import json
+
+def main(args):
+
+    '''
+    2022/06/17
+    Modify load_checkpoint to from_pretraind
+    '''
+    # initialize_distributed(args)
+    # load model from saved checkpoint 
+    
+    model_path = '/path/to/checkpoints/'
+    
+    model, args = CachedAutoregressiveModel.from_pretrained(args, model_path)
+
+    if args.fp16:
+        model = model.half()
+    model = model.to(args.device)
+    set_random_seed(args.seed)
+    model.eval()
+    
+    tokenizer = get_tokenizer(args)
+ 
+    # define function for each query
+    end_tokens = [tokenizer.get_command('eos').Id] 
+    strategy = BaseStrategy(temperature=args.temperature, top_k=args.top_k, end_tokens=end_tokens)
+    
+    def process(raw_text): 
+        if args.with_id:
+            query_id, raw_text = raw_text.split('\t')
+        raw_text = json.loads(raw_text)
+        question=raw_text["question"] + "答："
+        raw_text = question
+        seq = tokenizer._encode(raw_text)
+        if len(seq) != 0 and seq[0] == 20005:
+            seq = seq[1:]
+        seq = [tokenizer.get_command('ENC').Id] + seq
+        seq += [-1] * (args.max_sequence_length - len(seq)) 
+        if len(seq) > args.max_sequence_length:
+            raise ValueError('text too long.')
+        # generation
+        seq = torch.cuda.LongTensor(seq, device=args.device)
+        mbz = args.max_inference_batch_size
+        assert args.batch_size < mbz or args.batch_size % mbz == 0
+        output_list = []
+        for tim in range(max(args.batch_size // mbz, 1)):
+            output = filling_sequence(model, seq.clone(),
+                    batch_size=min(args.batch_size, mbz),
+                    strategy=strategy,
+                    log_attention_weights=None
+                    )[0]
+            if isinstance(output, torch.Tensor): # different strategies
+                output = list(output)
+
+            output_list.extend(output)
+        # find SEP to obatin output 
+        for i in range(len(output_list)):
+            output = output_list[i].tolist() 
+            try:
+                unfinished = output.index(-1)
+            except ValueError:
+                unfinished = len(output)
+            if output[unfinished - 1] in end_tokens:
+                unfinished -= 1
+            output_list[i] = output[1:unfinished]
+            bog = output.index(tokenizer.get_command('eos').Id)
+            output_list[i] = output[1:bog] + output[bog+1:unfinished] 
+                
+        # decoding 
+        txts = [] 
+        for seq in output_list:
+            decode_tokens = tokenizer.DecodeIds(seq)
+            txts.append(decode_tokens)
+        
+        # save
+        if args.with_id:
+            full_path = os.path.join(args.output_path, query_id + '.txt')
+        else:
+            prefix = raw_text.replace('/', '')[:20]
+            full_path = timed_name(prefix, '.txt', args.output_path)
+            print(txts[0]) # print the first.
+        test_eval_path = os.path.join(args.output_path, 'test_eval.txt')
+        with open(test_eval_path, 'a', encoding='utf-8') as fout:
+            fout.write(txts[0] + '\n')
+        os.chmod(test_eval_path, stat.S_IRWXO + stat.S_IRWXG + stat.S_IRWXU)
+
+    os.makedirs(args.output_path, exist_ok=True)
+    generate_continually(process, args.input_source) 
+
+
+if __name__ == "__main__":
+    py_parser = argparse.ArgumentParser(add_help=False)
+    
+    known, args_list = py_parser.parse_known_args()
+    args = get_args(args_list)
+    args = argparse.Namespace(**vars(args), **vars(known))
+    args.do_train = False
+    
+    with torch.no_grad():
+        main(args)
+