Updated code snippets for testing

README.md

@@ -6,6 +6,7 @@ tags:
   - solidity
   - web3
   - code generation
+  - smart contract
 widget:
 - text: "pragma solidity ^0.5.7;\n// Context: ParentA | Functions: helloA helloB | Constants: constantA \ncontract HelloWorld is ParentA {"
 ---
@@ -17,18 +18,52 @@ widget:
   - Contract/Library/Interface declaration header, e.g. `HelloWorld` ended with `{`
 
 ```python
+# !pip install transformers -q
+
 from transformers import AutoTokenizer, T5ForConditionalGeneration
+
+DEVICE = 'cuda'  # fallback to cpu if you do not have cuda
 tokenizer = AutoTokenizer.from_pretrained("hululuzhu/solidity-t5")
-model = T5ForConditionalGeneration.from_pretrained("hululuzhu/solidity-t5")
+model = T5ForConditionalGeneration.from_pretrained("hululuzhu/solidity-t5").to(DEVICE)
 
 text = """pragma solidity ^0.5.7;
 // Context: ParentA | Functions: helloA helloB | Constants: constantA 
 contract HelloWorld is ParentA {"""
-input_ids = model.tokenizer(text, return_tensors="pt", truncation=True).input_ids.to('cuda')
+input_ids = tokenizer(text, return_tensors="pt", truncation=True).input_ids.to(DEVICE)
 
 # Need to tune beam/topk/topp params to get good outcome
-generated_ids = model.model.generate(input_ids, max_length=256, num_beams=5, top_p=0.95, top_k=50)
-print(model.tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+generated_ids = model.generate(input_ids, max_length=256, num_beams=5, top_p=0.95, top_k=50)
+print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
+
+# Expect outcome
+"""
+string public constant name = "Hello World";
+string public constant symbol = "HELPER";
+uint8 public constant decimals = 18;
+uint256 public constant initialSupply = 0;
+uint256 public constant override returns (uint256) {
+return initialSupply;
+}
+function initialSupply() public view returns (uint256) {
+return initialSupply;
+}
+function balanceOf(address _owner) public view returns (uint256) {
+return balanceOf(_owner);
+}
+function transfer(address _to, uint256 _value) public returns (bool) {
+balanceOf[msg.sender] -= _value;
+balanceOf[_to] += _value;
+emit Transfer(msg.sender, _to, _value);
+return true;
+}
+function transferFrom(address _from, address _to, uint256 _value) public returns (bool) {
+balanceOf[_from] -= _value;
+balanceOf[_to] += _value;
+emit Transfer(_from, _to, _value);
+return true;
+}
+function approve(address _spender, uint256 _value) public returns (bool)
+"""
 ```