Update app.py

app.py

@@ -1,31 +1,72 @@
 import gradio as gr
 import torch
 import pickle
+import io
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
 
 # Load model and tokenizer from pickle files
 print("Loading model and tokenizer...")
 
+# Determine device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
 try:
-    # Load LoRA model
+    # Custom unpickler to handle device mapping
+    class CPU_Unpickler(pickle.Unpickler):
+        def find_class(self, module, name):
+            if module == 'torch.storage' and name == '_load_from_bytes':
+                return lambda b: torch.load(io.BytesIO(b), map_location=device)
+            else:
+                return super().find_class(module, name)
+    
+    # Load LoRA model with device mapping
+    print("Loading LoRA model...")
     with open('gpt2_pseudo2code_lora_model.pkl', 'rb') as f:
-        model = pickle.load(f)
+        if device == "cpu":
+            # Use custom unpickler for CPU
+            model = CPU_Unpickler(f).load()
+        else:
+            model = pickle.load(f)
     print("✓ Model loaded successfully")
     
     # Load tokenizer
+    print("Loading tokenizer...")
     with open('gpt2_pseudo2code_tokenizer.pkl', 'rb') as f:
         tokenizer = pickle.load(f)
     print("✓ Tokenizer loaded successfully")
     
-    # Move model to appropriate device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Ensure model is on correct device
     model = model.to(device)
     model.eval()
-    print(f"✓ Model moved to {device}")
+    print(f"✓ Model ready on {device}")
+    
+    # Print model info
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"✓ Total parameters: {total_params:,}")
+    print(f"✓ Trainable parameters: {trainable_params:,}")
     
 except Exception as e:
     print(f"Error loading model: {e}")
-    raise
+    print("\nTrying alternative loading method...")
+    try:
+        # Alternative method: load with torch.load directly
+        import io
+        
+        with open('gpt2_pseudo2code_lora_model.pkl', 'rb') as f:
+            buffer = io.BytesIO(f.read())
+            model = torch.load(buffer, map_location=torch.device(device))
+        
+        with open('gpt2_pseudo2code_tokenizer.pkl', 'rb') as f:
+            tokenizer = pickle.load(f)
+        
+        model = model.to(device)
+        model.eval()
+        print("✓ Model loaded successfully using alternative method")
+    except Exception as e2:
+        print(f"Alternative loading also failed: {e2}")
+        raise
 
 def generate_code(pseudocode, indent, line, max_length=128, temperature=0.7, top_p=0.9):
     """