Update app.py

app.py

@@ -1,70 +1,139 @@
+import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
 import bitsandbytes as bnb
+import gc
 
-model_name = "peterxyz/detect-llama-34b"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+@st.cache_resource
+def load_model():
+    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
+    )
+    
+    # Clear CUDA cache and garbage collect
+    torch.cuda.empty_cache()
+    gc.collect()
+    
+    model_nf4 = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        quantization_config=nf4_config,
+        device_map="auto"
+    )
+    model = PeftModel.from_pretrained(model_nf4, model_name)
+    
+    return model, tokenizer
+
+def analyze_contract(contract_code, model, tokenizer):
+    inputs = tokenizer(
+        f"{contract_code}\n\nidentify vulnerability of this code given above",
+        return_tensors="pt"
+    ).to("cuda")
+    
+    outputs = model.generate(
+        **inputs,
+        max_length=1024,
+        temperature=0.7,
+        num_return_sequences=1
+    )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Set page config
+st.set_page_config(
+    page_title="Smart Contract Vulnerability Detector",
+    page_icon="🔍",
+    layout="wide"
+)
 
-nf4_config = BitsAndBytesConfig(
-   load_in_4bit=True,
-   bnb_4bit_quant_type="nf4",
-   bnb_4bit_use_double_quant=True,
-   bnb_4bit_compute_dtype=torch.bfloat16
+# Main app
+st.title("🔍 Smart Contract Vulnerability Detector")
+st.markdown("""
+This app analyzes Solidity smart contracts for potential vulnerabilities using a fine-tuned LLaMA model.
+Simply paste your smart contract code below and click 'Analyze'.
+""")
+
+# Initialize session state for the model
+if 'model' not in st.session_state:
+    with st.spinner('Loading model... This might take a few minutes...'):
+        st.session_state.model, st.session_state.tokenizer = load_model()
+        st.success('Model loaded successfully!')
+
+# Create the main interface
+contract_code = st.text_area(
+    "Paste your Solidity contract code here:",
+    height=300,
+    placeholder="pragma solidity ^0.5.0;\n\ncontract YourContract {\n    // Your code here\n}"
 )
 
-import gc
-torch.cuda.empty_cache()
-gc.collect()
+analyze_button = st.button("Analyze Contract", type="primary")
 
-model_nf4 = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=nf4_config)
-model = PeftModel.from_pretrained(model_nf4, model_name)
+# Sample contract button
+if st.button("Load Sample Contract"):
+    contract_code = """pragma solidity ^0.5.0;
 
-from datasets import load_dataset
+contract ModifierEntrancy {
+    mapping (address => uint) public tokenBalance;
+    string constant name = "Nu Token";
+    Bank bank;
 
-dataset = load_dataset("peterxyz/smart-contract-vuln-detection")
+    constructor() public{
+        bank = new Bank();
+    }
 
-input_text = """
-pragma solidity ^0.5.0;
+    function airDrop() hasNoBalance supportsToken public{
+        tokenBalance[msg.sender] += 20;
+    }
 
-contract ModifierEntrancy {
+    modifier supportsToken() {
+        require(keccak256(abi.encodePacked("Nu Token")) == bank.supportsToken());
+        _;
+    }
 
-  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);
-      _;
-  }
+    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))
+}"""
+    st.session_state.contract_code = contract_code
+    st.experimental_rerun()
+
+# Analysis section
+if analyze_button and contract_code:
+    with st.spinner('Analyzing contract...'):
+        try:
+            analysis = analyze_contract(
+                contract_code,
+                st.session_state.model,
+                st.session_state.tokenizer
+            )
+            
+            st.subheader("Analysis Results")
+            st.markdown(analysis)
+            
+        except Exception as e:
+            st.error(f"An error occurred during analysis: {str(e)}")
+            
+elif analyze_button:
+    st.warning("Please enter some contract code to analyze.")
+
+# Add footer with information
+st.markdown("---")
+st.markdown("""
+<div style='text-align: center'>
+    <p>Built with Streamlit and Hugging Face Transformers</p>
+    <p>Model: peterxyz/detect-llama-34b</p>
+</div>
+""", unsafe_allow_html=True)