Update app.py

we need to explicitly use LlamaTokenizer since the model is Llama-based.

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, LlamaTokenizer, BitsAndBytesConfig
 from peft import PeftModel
 import bitsandbytes as bnb
 import gc
@@ -8,7 +8,8 @@ import gc
 @st.cache_resource
 def load_model():
     model_name = "peterxyz/detect-llama-34b"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # Use LlamaTokenizer instead of AutoTokenizer
+    tokenizer = LlamaTokenizer.from_pretrained(model_name)
     
     nf4_config = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -24,23 +25,35 @@ def load_model():
     model_nf4 = AutoModelForCausalLM.from_pretrained(
         model_name, 
         quantization_config=nf4_config,
-        device_map="auto"
+        device_map="auto",
+        trust_remote_code=True  # Added this parameter for safety
     )
     model = PeftModel.from_pretrained(model_nf4, model_name)
     
     return model, tokenizer
 
 def analyze_contract(contract_code, model, tokenizer):
+    prompt = f"{contract_code}\n\nidentify vulnerability of this code given above"
+    
+    # Add padding token if needed
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
     inputs = tokenizer(
-        f"{contract_code}\n\nidentify vulnerability of this code given above",
-        return_tensors="pt"
+        prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=2048  # Added max length for safety
     ).to("cuda")
     
     outputs = model.generate(
         **inputs,
         max_length=1024,
         temperature=0.7,
-        num_return_sequences=1
+        num_return_sequences=1,
+        pad_token_id=tokenizer.pad_token_id,  # Explicitly set pad token ID
+        eos_token_id=tokenizer.eos_token_id   # Explicitly set EOS token ID
     )
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -59,11 +72,20 @@ This app analyzes Solidity smart contracts for potential vulnerabilities using a
 Simply paste your smart contract code below and click 'Analyze'.
 """)
 
+# Add a loading message while initializing
+if 'model_loaded' not in st.session_state:
+    st.session_state.model_loaded = False
+
 # 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!')
+if not st.session_state.model_loaded:
+    try:
+        with st.spinner('Loading model... This might take a few minutes...'):
+            st.session_state.model, st.session_state.tokenizer = load_model()
+            st.session_state.model_loaded = True
+            st.success('Model loaded successfully!')
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        st.stop()
 
 # Create the main interface
 contract_code = st.text_area(
@@ -72,10 +94,14 @@ contract_code = st.text_area(
     placeholder="pragma solidity ^0.5.0;\n\ncontract YourContract {\n    // Your code here\n}"
 )
 
-analyze_button = st.button("Analyze Contract", type="primary")
+col1, col2 = st.columns([1, 4])
+with col1:
+    analyze_button = st.button("Analyze Contract", type="primary")
+with col2:
+    load_sample = st.button("Load Sample Contract")
 
 # Sample contract button
-if st.button("Load Sample Contract"):
+if load_sample:
     contract_code = """pragma solidity ^0.5.0;
 
 contract ModifierEntrancy {
@@ -112,8 +138,8 @@ contract Bank{
 
 # Analysis section
 if analyze_button and contract_code:
-    with st.spinner('Analyzing contract...'):
-        try:
+    try:
+        with st.spinner('Analyzing contract...'):
             analysis = analyze_contract(
                 contract_code,
                 st.session_state.model,
@@ -121,10 +147,15 @@ if analyze_button and contract_code:
             )
             
             st.subheader("Analysis Results")
-            st.markdown(analysis)
             
-        except Exception as e:
-            st.error(f"An error occurred during analysis: {str(e)}")
+            # Create an expandable section for the analysis
+            with st.expander("View Full Analysis", expanded=True):
+                st.markdown(analysis)
+            
+    except Exception as e:
+        st.error(f"An error occurred during analysis: {str(e)}")
+        st.markdown("**Debug Information:**")
+        st.code(str(e))
             
 elif analyze_button:
     st.warning("Please enter some contract code to analyze.")