Update app.py

app.py

@@ -17,19 +17,18 @@ logging.basicConfig(
 def load_model():
     """
     Loads and caches the pre-trained language model and tokenizer.
-
     Returns:
         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 = AutoModelForCausalLM.from_pretrained(
             model_path,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto",
-            load_in_4bit=False,
-            load_in_8bit=True,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            device_map={"": device},  # This will specify CPU or GPU explicitly
+            load_in_8bit=False,  # Disabled for stability
             trust_remote_code=True,
         )
         tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
@@ -46,7 +45,6 @@ def sanitize_input(text):
     
     Args:
         text (str): User input text.
-
     Returns:
         str: Sanitized text.
     """
@@ -59,19 +57,18 @@ def sanitize_input(text):
 def generate_text(model, tokenizer, instruction):
     """
     Generates text based on the provided instruction using the loaded model.
-
     Args:
         model: The language model.
         tokenizer: Tokenizer for encoding/decoding.
         instruction (str): Instruction text for the model.
-
     Returns:
         str: Generated text response from the model.
     """
     try:
         # Validate and sanitize instruction input
         instruction = sanitize_input(instruction)
-        tokens = tokenizer.encode(instruction, return_tensors='pt').to('cuda')
+        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, 
@@ -90,7 +87,6 @@ def generate_text(model, tokenizer, instruction):
 def load_json_data():
     """
     Loads JSON data, simulating the loading process with a sample list.
-
     Returns:
         list: A list of dictionaries with sample user data.
     """
@@ -138,4 +134,3 @@ for user in user_data:
     st.write(f"**Country:** {user['country']}")
     st.write(f"**Company:** {user['company']}")
     st.write("---")
-