init commit

README.md

@@ -46,18 +46,18 @@ import torch
 from transformers import T5Tokenizer, T5Config, T5ForConditionalGeneration
 
 # load tokenizer and model 
-pretrained_model = "IDEA-CCNL/Randeng-T5-784M-MultiTask-Chinese"
+pretrained_model = "/cognitive_comp/wuxiaojun/pretrained/pytorch/huggingface/Randeng-T5-784M-MultiTask-Chinese"
 
 special_tokens = ["<extra_id_{}>".format(i) for i in range(100)]
 tokenizer = T5Tokenizer.from_pretrained(
-    args.pretrained_model,
+    pretrained_model,
     do_lower_case=True,
     max_length=512,
     truncation=True,
     additional_special_tokens=special_tokens,
 )
-config = T5Config.from_pretrained(args.pretrained_model)
-model = T5ForConditionalGeneration.from_pretrained(args.pretrained_model, config=config)
+config = T5Config.from_pretrained(pretrained_model)
+model = T5ForConditionalGeneration.from_pretrained(pretrained_model, config=config)
 model.resize_token_embeddings(len(tokenizer))
 model.eval()
 
@@ -66,8 +66,8 @@ text = "新闻分类任务：【微软披露拓扑量子计算机计划！】这
 encode_dict = tokenizer(text, max_length=512, padding='max_length',truncation=True)
 
 inputs = {
-  "input_ids": torch.tensor(encode_dict['input_ids']).long(),
-  "attention_mask": torch.tensor(encode_dict['attention_mask']).long(),
+  "input_ids": torch.tensor([encode_dict['input_ids']]).long(),
+  "attention_mask": torch.tensor([encode_dict['attention_mask']]).long(),
   }
 
 # generate answer
@@ -80,8 +80,9 @@ logits = model.generate(
 
 logits=logits[:,1:]
 predict_label = [tokenizer.decode(i,skip_special_tokens=True) for i in logits]
+print(predict_label)
 
-# model Output: 科技
+# model output: 科技
 ```
 
 ## 引用 Citation