Update app.py

app.py

@@ -1,71 +1,105 @@
 import gradio as gr
 import torch
-import pickle
-import io
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from transformers import GPT2Tokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig, AutoPeftModelForCausalLM
+import os
 
-# Load model and tokenizer from pickle files
-print("Loading model and tokenizer...")
+# Load model and tokenizer
+print("="*70)
+print("Loading Pseudo-Code to Code Generator")
+print("="*70)
 
 # Determine device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
 try:
-    # 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)
+    # Method 1: Try loading as PeftModel (LoRA adapters)
+    print("\nAttempting to load LoRA model from Hugging Face format...")
+    model_path = "./model"  # or wherever you uploaded the model files
     
-    # Load LoRA model with device mapping
-    print("Loading LoRA model...")
-    with open('gpt2_pseudo2code_lora_model.pkl', 'rb') as 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")
+    # Check if model files exist
+    if os.path.exists(model_path):
+        print(f"✓ Model directory found: {model_path}")
+        
+        # Load with AutoPeftModel (handles LoRA automatically)
+        model = AutoPeftModelForCausalLM.from_pretrained(
+            model_path,
+            device_map={"": device},
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        )
+        tokenizer = GPT2Tokenizer.from_pretrained(model_path)
+        
+        print("✓ Model and tokenizer loaded successfully (Hugging Face format)")
+    else:
+        # Fallback: Load from current directory
+        print(f"✗ Model directory not found, trying current directory...")
+        model = AutoPeftModelForCausalLM.from_pretrained(
+            ".",
+            device_map={"": device},
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        )
+        tokenizer = GPT2Tokenizer.from_pretrained(".")
+        print("✓ Model loaded from current directory")
     
-    # Ensure model is on correct device
-    model = model.to(device)
+    # Set model to evaluation mode
     model.eval()
     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:,}")
+    try:
+        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:
+        print("✓ Model parameters info not available")
+    
+    print("="*70)
     
 except Exception as e:
-    print(f"Error loading model: {e}")
-    print("\nTrying alternative loading method...")
+    print(f"\n✗ Error loading with AutoPeftModel: {e}")
+    print("\nTrying alternative method: Loading base model + LoRA adapters separately...")
+    
     try:
-        # Alternative method: load with torch.load directly
-        import io
+        # Method 2: Load base GPT-2 and merge LoRA adapters
+        from transformers import GPT2LMHeadModel
+        
+        print("Loading base GPT-2 model...")
+        base_model = GPT2LMHeadModel.from_pretrained("gpt2")
         
-        with open('gpt2_pseudo2code_lora_model.pkl', 'rb') as f:
-            buffer = io.BytesIO(f.read())
-            model = torch.load(buffer, map_location=torch.device(device))
+        print("Loading LoRA adapters...")
+        model = PeftModel.from_pretrained(
+            base_model,
+            model_path if os.path.exists(model_path) else ".",
+            device_map={"": device}
+        )
         
-        with open('gpt2_pseudo2code_tokenizer.pkl', 'rb') as f:
-            tokenizer = pickle.load(f)
+        tokenizer = GPT2Tokenizer.from_pretrained(
+            model_path if os.path.exists(model_path) else "."
+        )
         
-        model = model.to(device)
         model.eval()
-        print("✓ Model loaded successfully using alternative method")
+        print("✓ Model loaded successfully (base + adapters)")
+        
     except Exception as e2:
-        print(f"Alternative loading also failed: {e2}")
+        print(f"\n✗ Alternative method also failed: {e2}")
+        print("\n" + "="*70)
+        print("DEPLOYMENT INSTRUCTIONS")
+        print("="*70)
+        print("Please upload the model in Hugging Face format, not pickle!")
+        print("\nFiles needed:")
+        print("  - adapter_config.json")
+        print("  - adapter_model.safetensors (or .bin)")
+        print("  - tokenizer.json")
+        print("  - tokenizer_config.json")
+        print("  - special_tokens_map.json")
+        print("  - vocab.json")
+        print("  - merges.txt")
+        print("\nSee SAVE_MODEL_FOR_HF.py for instructions on how to save properly.")
+        print("="*70)
         raise
 
 def generate_code(pseudocode, indent, line, max_length=128, temperature=0.7, top_p=0.9):