Update app.py

app.py

@@ -1,38 +1,55 @@
 import streamlit as st
 import torch
-from transformers import AutoModelForCausalLM, LlamaTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, LlamaTokenizer
 from peft import PeftModel
-import bitsandbytes as bnb
 import gc
 
 @st.cache_resource
 def load_model():
     model_name = "peterxyz/detect-llama-34b"
-    # Use LlamaTokenizer instead of AutoTokenizer
     tokenizer = LlamaTokenizer.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
-    )
+    # Check if CUDA is available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    st.info(f"Using device: {device}")
     
-    # Clear CUDA cache and garbage collect
-    torch.cuda.empty_cache()
+    # Clear memory
+    if device == "cuda":
+        torch.cuda.empty_cache()
     gc.collect()
     
-    model_nf4 = AutoModelForCausalLM.from_pretrained(
-        model_name, 
-        quantization_config=nf4_config,
-        device_map="auto",
-        trust_remote_code=True  # Added this parameter for safety
-    )
-    model = PeftModel.from_pretrained(model_nf4, model_name)
+    # Load model with appropriate settings based on device
+    if device == "cuda":
+        from transformers import BitsAndBytesConfig
+        import bitsandbytes as bnb
+        
+        nf4_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        
+        model_nf4 = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=nf4_config,
+            device_map="auto",
+            trust_remote_code=True
+        )
+        model = PeftModel.from_pretrained(model_nf4, model_name)
+    else:
+        # For CPU, load with reduced precision but without 4-bit quantization
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,  # Use float32 for CPU
+            device_map={"": device},
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
+        )
     
-    return model, tokenizer
+    return model, tokenizer, device
 
-def analyze_contract(contract_code, model, tokenizer):
+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
@@ -44,16 +61,16 @@ def analyze_contract(contract_code, model, tokenizer):
         return_tensors="pt",
         padding=True,
         truncation=True,
-        max_length=2048  # Added max length for safety
-    ).to("cuda")
+        max_length=2048
+    ).to(device)
     
     outputs = model.generate(
         **inputs,
         max_length=1024,
         temperature=0.7,
         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
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id
     )
     
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -80,7 +97,7 @@ if 'model_loaded' not in st.session_state:
 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, st.session_state.tokenizer, st.session_state.device = load_model()
             st.session_state.model_loaded = True
             st.success('Model loaded successfully!')
     except Exception as e:
@@ -143,7 +160,8 @@ if analyze_button and contract_code:
             analysis = analyze_contract(
                 contract_code,
                 st.session_state.model,
-                st.session_state.tokenizer
+                st.session_state.tokenizer,
+                st.session_state.device
             )
             
             st.subheader("Analysis Results")