Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3model_type = 'kdf/python-docstring-generation'
4tokenizer = AutoTokenizer.from_pretrained(model_type)
5model = AutoModelForCausalLM.from_pretrained(model_type)
6
7inputs = tokenizer('''<|endoftext|>
8def load_excel(path):
9 return pd.read_excel(path)
10
11# docstring
12"""''', return_tensors='pt')
13
14doc_max_length = 128
15
16generated_ids = model.generate(
17 **inputs,
18 max_length=inputs.input_ids.shape[1] + doc_max_length,
19 do_sample=False,
20 return_dict_in_generate=True,
21 num_return_sequences=1,
22 output_scores=True,
23 pad_token_id=50256,
24 eos_token_id=50256 # <|endoftext|>
25)
26
27ret = tokenizer.decode(generated_ids.sequences[0], skip_special_tokens=False)
28print(ret)
291inputs = tokenizer('''<|endoftext|>
2def add(a, b):
3 return a + b
4
5# docstring
6"""
7 Calculate numbers add.
8
9 Args:
10 a: the first number to add
11 b: the second number to add
12
13 Return:
14 The result of a + b
15"""
16<|endoftext|>
17def load_excel(path):
18 return pd.read_excel(path)
19
20# docstring
21"""''', return_tensors='pt')
22
23doc_max_length = 128
24
25generated_ids = model.generate(
26 **inputs,
27 max_length=inputs.input_ids.shape[1] + doc_max_length,
28 do_sample=False,
29 return_dict_in_generate=True,
30 num_return_sequences=1,
31 output_scores=True,
32 pad_token_id=50256,
33 eos_token_id=50256 # <|endoftext|>
34)
35
36ret = tokenizer.decode(generated_ids.sequences[0], skip_special_tokens=False)
37print(ret)
38
39inputs = tokenizer('''<|endoftext|>
40def add(a, b):
41 return a + b
42
43# docstring
44"""
45 计算数字相加
46
47 Args:
48 a: 第一个加数
49 b: 第二个加数
50
51 Return:
52 相加的结果
53"""
54<|endoftext|>
55def load_excel(path):
56 return pd.read_excel(path)
57
58# docstring
59"""''', return_tensors='pt')
60
61doc_max_length = 128
62
63generated_ids = model.generate(
64 **inputs,
65 max_length=inputs.input_ids.shape[1] + doc_max_length,
66 do_sample=False,
67 return_dict_in_generate=True,
68 num_return_sequences=1,
69 output_scores=True,
70 pad_token_id=50256,
71 eos_token_id=50256 # <|endoftext|>
72)
73
74ret = tokenizer.decode(generated_ids.sequences[0], skip_special_tokens=False)
75print(ret)