Update app.py

app.py

@@ -21,9 +21,9 @@ def load_model():
         model: Pre-trained language model.
         tokenizer: Tokenizer for the model.
     """
-    model_path = "Canstralian/pentest_ai"
     try:
         device = "cuda" if torch.cuda.is_available() else "cpu"
+        model_path = "Canstralian/pentest_ai"  # Replace with the actual path if different
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
             torch_dtype=torch.float16 if device == "cuda" else torch.float32,
@@ -70,10 +70,10 @@ def generate_text(model, tokenizer, instruction):
         device = "cuda" if torch.cuda.is_available() else "cpu"
         tokens = tokenizer.encode(instruction, return_tensors='pt').to(device)
         generated_tokens = model.generate(
-            tokens, 
-            max_length=1024, 
-            top_p=1.0, 
-            temperature=0.5, 
+            tokens,
+            max_length=1024,
+            top_p=1.0,
+            temperature=0.5,
             top_k=50
         )
         generated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
@@ -109,6 +109,9 @@ st.title("Penetration Testing AI Assistant")
 # Load the model and tokenizer
 model, tokenizer = load_model()
 
+if not model or not tokenizer:
+    st.error("Failed to load model or tokenizer. Please check your configuration.")
+
 # User instruction input
 instruction = st.text_input("Enter an instruction for the model:")