lewiswu1209 commited on
Commit
b56b764
1 Parent(s): c1e6869

add fill_blanks.py

Browse files
Files changed (1) hide show
  1. bot/skills/fill_blanks.py +46 -0
bot/skills/fill_blanks.py ADDED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import re
4
+ import json
5
+ import requests
6
+
7
+ class Skill:
8
+ def __init__(self:object) -> None:
9
+ pass
10
+
11
+ def set_tokenizer(self:object, tokenizer:object):
12
+ self.tokenizer = tokenizer
13
+
14
+ def predict(self, question:str, seed=1234, out_seq_length=256, min_gen_length=1, sampling_strategy="BaseStrategy", num_beams=2, length_penalty=1, no_repeat_ngram_size=3, temperature=0.1, topk=1, topp=0.1):
15
+ url = 'https://pretrain.aminer.cn/api/v2/completions'
16
+ prompt = question
17
+ payload = json.dumps({
18
+ "apikey": os.environ.get("WINNIE_APIKEY"),
19
+ "apisecret": os.environ.get("WINNIE_APISECRET"),
20
+ "model": "glm",
21
+ "language": "zh-CN",
22
+ "prompt": prompt,
23
+ "temperature": temperature,
24
+ "top_k": topk,
25
+ "top_p": topp,
26
+ "max_tokens": out_seq_length,
27
+ "stop": ["\n"],
28
+ "presence_penalty": 2,
29
+ "frequency_penalty": 2
30
+ })
31
+
32
+ headers = {
33
+ 'Content-Type': 'application/json'
34
+ }
35
+ response = requests.request("POST", url, headers=headers, data=payload).json()
36
+ if "output" in response["result"]:
37
+ answer = response["result"]["output"]["raw"].split("<|startofpiece|>")[-1]
38
+ return answer
39
+
40
+ def process(self:object, input_txt:str, history_list:list, role_card:dict):
41
+ output_text:str = None
42
+ if input_txt.find("[MASK]") != -1:
43
+ history_list.append( self.tokenizer.encode(input_txt, add_special_tokens=False) )
44
+ output_text = self.predict(input_txt)
45
+ history_list.append( self.tokenizer.encode(output_text, add_special_tokens=False) )
46
+ return output_text, history_list, role_card