Update app.py

app.py

@@ -49,7 +49,7 @@ def load_model():
     
     return model, tokenizer, device
 
-def analyze_contract(contract_code, model, tokenizer, device, batch_size=1):
+def analyze_contract(contract_code, model, tokenizer, device):
     prompt = f"{contract_code}\n\nidentify vulnerability of this code given above"
     
     # Add padding token if needed
@@ -61,8 +61,7 @@ def analyze_contract(contract_code, model, tokenizer, device, batch_size=1):
         return_tensors="pt",
         padding=True,
         truncation=True,
-        max_length=2048,
-        batch_size=batch_size
+        max_length=2048
     ).to(device)
     
     outputs = model.generate(
@@ -112,13 +111,11 @@ contract_code = st.text_area(
     placeholder="pragma solidity ^0.5.0;\n\ncontract YourContract {\n    // Your code here\n}"
 )
 
-col1, col2, col3 = st.columns([1, 2, 1])
+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")
-with col3:
-    batch_size = st.number_input("Batch Size", min_value=1, max_value=8, value=1, step=1)
 
 # Sample contract button
 if load_sample:
@@ -160,17 +157,11 @@ contract Bank{
 if analyze_button and contract_code:
     try:
         with st.spinner('Analyzing contract...'):
-            # Get the current available GPU memory
-            if st.session_state.device == "cuda":
-                gpu_stats = torch.cuda.get_device_properties(0)
-                st.write(f"GPU Memory: {gpu_stats.total_memory / (1024 ** 2):.2f} MB")
-            
             analysis = analyze_contract(
                 contract_code,
                 st.session_state.model,
                 st.session_state.tokenizer,
-                st.session_state.device,
-                batch_size=batch_size
+                st.session_state.device
             )
             
             st.subheader("Analysis Results")
@@ -187,10 +178,6 @@ if analyze_button and contract_code:
 elif analyze_button:
     st.warning("Please enter some contract code to analyze.")
 
-# Display a warning about CPU performance
-if st.session_state.device == "cpu":
-    st.warning("Running the model on CPU can be significantly slower. For better performance, consider using a GPU-enabled environment.")
-
 # Add footer with information
 st.markdown("---")
 st.markdown("""