Update code example

README.md

@@ -43,57 +43,13 @@ from transformers import RobertaTokenizer, T5ForConditionalGeneration
 tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-small')
 model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-small')
 
-text = "def greet(user): print(f'hello <extra_id_0>!') </s>"
-inputs = tokenizer(text, return_tensors="pt").input_ids
+text = "def greet(user): print(f'hello <extra_id_0>!')"
+input_ids = tokenizer(text, return_tensors="pt").input_ids
 
 # simply generate a single sequence
-generated_ids = model.generate(input_ids, max_length=8)
+generated_ids = model.generate(input_ids, max_length=10)
 print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
-# this prints {user.name}
-
-# or, generating 20 sequences with maximum length set to 10
-outputs = model.generate(input_ids=input_ids, 
-                          num_beams=200, num_return_sequences=20,
-                          max_length=10)
-
-_0_index = text.index('<extra_id_0>')
-_result_prefix = text[:_0_index]
-_result_suffix = text[_0_index+12:]  # 12 is the length of <extra_id_0>
-
-def _filter(output, end_token='<extra_id_1>'):
-    # The first token is <pad> (indexed at 0), the second token is <s> (indexed at 1)
-    # and the third token is <extra_id_0> (indexed at 32099)
-    # So we only decode from the fourth generated id
-    _txt = tokenizer.decode(output[3:], skip_special_tokens=False, clean_up_tokenization_spaces=False)
-    if end_token in _txt:
-        _end_token_index = _txt.index(end_token)
-        return _result_prefix + _txt[:_end_token_index] + _result_suffix
-    else:
-        return _result_prefix + _txt + _result_suffix
-
-results = list(map(_filter, outputs))
-print(results)
-# this prints:
-#["def greet(user): print(f'hello {user.name} {user!') </s>",
-# "def greet(user): print(f'hello {user.username} {user!') </s>",
-# "def greet(user): print(f'hello {user.name}: {user!') </s>",
-# "def greet(user): print(f'hello {user}') print(f!') </s>",
-# "def greet(user): print(f'hello {user.name} �!') </s>",
-# "def greet(user): print(f'hello {user}') print ( f!') </s>",
-# "def greet(user): print(f'hello {user.username}: {user!') </s>",
-# "def greet(user): print(f'hello {user}' ) print(f!') </s>",
-# "def greet(user): print(f'hello {user.username} �!') </s>",
-# "def greet(user): print(f'hello {user.name}, {user!') </s>",
-# "def greet(user): print(f'hello {user.login} {user!') </s>",
-# "def greet(user): print(f'hello {user} →!') </s>",
-# "def greet(user): print(f'hello {user}!') print(!') </s>",
-# "def greet(user): print(f'hello {user.name} ({user!') </s>",
-# "def greet(user): print(f'hello {user.email} {user!') </s>",
-# "def greet(user): print(f'hello {user}!') print (!') </s>",
-# "def greet(user): print(f'hello {user.username}, {user!') </s>",
-# "def greet(user): print(f'hello {user}' ) print ( f!') </s>",
-# "def greet(user): print(f'hello {user.nickname} {!') </s>",
-# "def greet(user): print(f'hello {user} {user.name!') </s>"]
+# this prints "user: {user.name}"
 ```
 
 ## Training data