Update app.py

app.py

@@ -0,0 +1,70 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftModel
+import bitsandbytes as bnb
+
+model_name = "peterxyz/detect-llama-34b"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+nf4_config = BitsAndBytesConfig(
+   load_in_4bit=True,
+   bnb_4bit_quant_type="nf4",
+   bnb_4bit_use_double_quant=True,
+   bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+import gc
+torch.cuda.empty_cache()
+gc.collect()
+
+model_nf4 = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config)
+model = PeftModel.from_pretrained(model_nf4, model_name)
+
+from datasets import load_dataset
+
+dataset = load_dataset("peterxyz/smart-contract-vuln-detection")
+
+input_text = """
+pragma solidity ^0.5.0;
+
+contract ModifierEntrancy {
+
+  mapping (address => uint) public tokenBalance;
+  string constant name = "Nu Token";
+  Bank bank;
+
+  constructor() public{
+      bank = new Bank();
+  }
+
+  //If a contract has a zero balance and supports the token give them some token
+  function airDrop() hasNoBalance supportsToken  public{
+    tokenBalance[msg.sender] += 20;
+  }
+
+  //Checks that the contract responds the way we want
+  modifier supportsToken() {
+    require(keccak256(abi.encodePacked("Nu Token")) == bank.supportsToken());
+    _;
+  }
+
+  //Checks that the caller has a zero balance
+  modifier hasNoBalance {
+      require(tokenBalance[msg.sender] == 0);
+      _;
+  }
+}
+
+contract Bank{
+
+    function supportsToken() external returns(bytes32) {
+        return keccak256(abi.encodePacked("Nu Token"));
+    }
+
+}
+
+identify vulnerability of this code given above
+"""
+inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
+outputs = model.generate(**inputs)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))