.

app.py

@@ -45,20 +45,47 @@ try:
     
 except Exception as e:
     print(f"Error initializing model {MODEL_NAME}: {str(e)}")
-    print("Falling back to tiny-gpt2...")
-    MODEL_NAME = "sshleifer/tiny-gpt2"
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    text_pipeline = pipeline(
-        "text-generation",
-        model=MODEL_NAME,
-        tokenizer=tokenizer,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        trust_remote_code=True
-    )
-    model = text_pipeline.model
-    tok = text_pipeline.tokenizer
-    print(f"Fallback model loaded: {MODEL_NAME}")
+    print("Trying with simplified parameters...")
+    
+    try:
+        # Try with simpler parameters
+        text_pipeline = pipeline(
+            "text-generation",
+            model=MODEL_NAME,
+            trust_remote_code=True
+        )
+        model = text_pipeline.model
+        tok = text_pipeline.tokenizer
+        print(f"Model loaded with simplified parameters: {MODEL_NAME}")
+        
+    except Exception as e2:
+        print(f"Second attempt failed: {str(e2)}")
+        print("Falling back to distilgpt2 (uses safetensors)...")
+        
+        # Use distilgpt2 which uses safetensors format and is more compatible
+        MODEL_NAME = "distilgpt2"
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+            text_pipeline = pipeline(
+                "text-generation",
+                model=MODEL_NAME,
+                tokenizer=tokenizer
+            )
+            model = text_pipeline.model
+            tok = text_pipeline.tokenizer
+            print(f"Fallback model loaded: {MODEL_NAME}")
+            
+        except Exception as e3:
+            print(f"Fallback also failed: {str(e3)}")
+            print("Trying direct model loading as last resort...")
+            
+            # Last resort: direct loading without pipeline
+            try:
+                tok = AutoTokenizer.from_pretrained(MODEL_NAME)
+                model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).eval()
+                print(f"Direct loading successful: {MODEL_NAME}")
+            except Exception as e4:
+                raise RuntimeError(f"All model loading attempts failed. Last error: {str(e4)}")
 
 # Log device information
 if hasattr(model, 'device'):

requirements.txt

@@ -1,6 +1,6 @@
 # Core dependencies for LLM Threat Association Analysis (ZeroGPU compatible)
 gradio>=4.0.0
-torch==2.4.0
+torch==2.5.1
 transformers>=4.30.0
 pandas>=2.0.0
 accelerate>=0.26.0