Update app.py

app.py

@@ -1,321 +1,253 @@
 import gradio as gr
 import torch
-from transformers import GPT2Tokenizer, AutoModelForCausalLM
-from peft import PeftModel, PeftConfig, AutoPeftModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
-# 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}")
+# Model configuration
+MODEL_NAME = "your-username/your-model-name"  # Replace with your actual HF model repo
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-try:
-    # 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 model and tokenizer
+@gr.utils.cache
+def load_model():
+    """Load the model and tokenizer with caching"""
+    print(f"Loading model from: {MODEL_NAME}")
+    print(f"Using device: {DEVICE}")
     
-    # 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")
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     
-    # Set model to evaluation mode
-    model.eval()
-    print(f"✓ Model ready on {device}")
+    # Set pad token if not set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     
-    # Print model info
-    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)
+    # Load model with appropriate settings
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+        device_map="auto" if DEVICE == "cuda" else None,
+        trust_remote_code=True
+    )
     
-except Exception as e:
-    print(f"\n✗ Error loading with AutoPeftModel: {e}")
-    print("\nTrying alternative method: Loading base model + LoRA adapters separately...")
+    if DEVICE == "cpu":
+        model = model.to(DEVICE)
     
-    try:
-        # 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")
-        
-        print("Loading LoRA adapters...")
-        model = PeftModel.from_pretrained(
-            base_model,
-            model_path if os.path.exists(model_path) else ".",
-            device_map={"": device}
-        )
-        
-        tokenizer = GPT2Tokenizer.from_pretrained(
-            model_path if os.path.exists(model_path) else "."
-        )
-        
-        model.eval()
-        print("✓ Model loaded successfully (base + adapters)")
-        
-    except Exception as 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
+    print("✅ Model and tokenizer loaded successfully!")
+    return model, tokenizer
+
+# Initialize model and tokenizer
+model, tokenizer = load_model()
 
-def generate_code(pseudocode, indent, line, max_length=128, temperature=0.7, top_p=0.9):
+def generate_code(pseudocode, indent=1, line=1, temperature=0.7, top_p=0.9, max_length=128):
     """
     Generate code from pseudo-code with line and indent information.
     
     Args:
         pseudocode: Input pseudo-code string
-        indent: Indentation level
-        line: Line number
-        max_length: Maximum length of generated sequence
-        temperature: Sampling temperature
-        top_p: Nucleus sampling parameter
+        indent: Indentation level (1-10)
+        line: Line number (1-100)
+        temperature: Sampling temperature (0.1-2.0)
+        top_p: Nucleus sampling parameter (0.1-1.0)
+        max_length: Maximum length of generated sequence (50-512)
     
     Returns:
         Generated code string
     """
     try:
-        # Format input with line and indent information
-        prompt = f"Pseudocode: {pseudocode} | Indent: {indent} | Line: {line}\nCode:"
+        # Validate inputs
+        if not pseudocode.strip():
+            return "❌ Error: Please enter some pseudocode."
         
-        # Tokenize input
-        inputs = tokenizer(prompt, return_tensors='pt', padding=True)
+        # Format input with line and indent information (matches training format)
+        prompt = f"Pseudocode: {pseudocode.strip()} | Indent: {indent} | Line: {line}\nCode:"
         
-        # Move to same device as model
-        device = next(model.parameters()).device
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors='pt', padding=True, truncation=True, max_length=256)
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
         
-        # Generate
+        # Generate with the model
         model.eval()
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_length=max_length,
-                temperature=temperature,
-                top_p=top_p,
+                max_new_tokens=max_length,
+                temperature=max(0.1, temperature),  # Ensure minimum temperature
+                top_p=max(0.1, top_p),  # Ensure minimum top_p
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
                 eos_token_id=tokenizer.eos_token_id,
-                num_return_sequences=1
+                num_return_sequences=1,
+                repetition_penalty=1.1,
+                no_repeat_ngram_size=2
             )
         
         # Decode output
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract only the code part
+        # Extract only the code part (remove the prompt)
         if "Code:" in generated_text:
-            code = generated_text.split("Code:")[1].strip()
+            code = generated_text.split("Code:")[-1].strip()
         else:
             code = generated_text.strip()
         
-        return code
-    
-    except Exception as e:
-        return f"Error generating code: {str(e)}"
-
-def gradio_generate_code(pseudocode, indent, line, temperature=0.7, top_p=0.9, max_length=128):
-    """
-    Wrapper function for Gradio interface.
-    """
-    if not pseudocode.strip():
-        return "⚠️ Please enter some pseudocode!"
-    
-    try:
-        indent = int(indent)
-        line = int(line)
-        generated_code = generate_code(
-            pseudocode,
-            indent,
-            line,
-            max_length=int(max_length),
-            temperature=float(temperature),
-            top_p=float(top_p)
-        )
-        return generated_code
-    except ValueError:
-        return "⚠️ Indent and Line must be valid numbers!"
+        # Clean up the output
+        if code.startswith(prompt):
+            code = code[len(prompt):].strip()
+        
+        return code if code else "❌ No code generated. Try adjusting the parameters."
+        
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        return f"❌ Error generating code: {str(e)}"
 
-# Example pseudocodes
-examples = [
-    ["create integer n", 1, 1, 0.7, 0.9, 128],
-    ["read n", 1, 2, 0.7, 0.9, 128],
-    ["for i from 0 to n", 1, 3, 0.7, 0.9, 128],
-    ["print i", 2, 4, 0.7, 0.9, 128],
-    ["if n is equal to 0", 1, 5, 0.7, 0.9, 128],
-    ["create string s", 1, 1, 0.7, 0.9, 128],
-    ["read s", 1, 2, 0.7, 0.9, 128],
-]
+def create_examples():
+    """Create example inputs for the interface"""
+    return [
+        ["create string s", 1, 1, 0.7, 0.9, 100],
+        ["read input from user", 1, 2, 0.7, 0.9, 100],
+        ["if s is empty", 1, 3, 0.7, 0.9, 100],
+        ["print hello world", 2, 4, 0.7, 0.9, 100],
+        ["for i from 0 to n", 1, 5, 0.7, 0.9, 100],
+        ["declare integer array", 1, 1, 0.5, 0.9, 80],
+        ["while condition is true", 2, 10, 0.8, 0.95, 120]
+    ]
 
 # Create Gradio interface
-with gr.Blocks(theme=gr.themes.Soft(), title="Pseudo-Code to Code Generator") as demo:
-    gr.Markdown(
-        """
-        # 🐍 Pseudo-Code to Code Generator (GPT-2 + LoRA)
-        
-        Convert natural language pseudo-code to executable code using a fine-tuned GPT-2 model with LoRA.
-        
-        **Model Details:**
-        - Base Model: GPT-2
-        - Training: SPOC Dataset (C++ code examples)
-        - Optimization: LoRA (Low-Rank Adaptation) + 16-bit precision
-        - Trained on: 20,000 pseudo-code to code pairs
-        
-        **Note:** The model was trained on C++ code examples from the SPOC dataset, so it generates C++-style code.
-        """
-    )
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    title="🐍 Pseudo-Code to Code Generator",
+    css="""
+    .gradio-container {
+        max-width: 1200px;
+        margin: auto;
+    }
+    .header {
+        text-align: center;
+        margin-bottom: 30px;
+    }
+    .info-box {
+        background-color: #f0f8ff;
+        padding: 15px;
+        border-radius: 10px;
+        margin: 10px 0;
+    }
+    """
+) as demo:
+    
+    # Header
+    gr.HTML("""
+    <div class="header">
+        <h1>🐍 Pseudo-Code to Code Generator</h1>
+        <p>Convert natural language pseudo-code to executable code using fine-tuned GPT-2</p>
+    </div>
+    """)
+    
+    # Info box
+    gr.HTML("""
+    <div class="info-box">
+        <h3>📋 How to use:</h3>
+        <ol>
+            <li><strong>Enter pseudocode:</strong> Describe what you want the code to do in natural language</li>
+            <li><strong>Set context:</strong> Adjust indent level and line number for better structure</li>
+            <li><strong>Tune generation:</strong> Modify temperature and top_p for different creativity levels</li>
+            <li><strong>Generate:</strong> Click submit to get your code!</li>
+        </ol>
+        <p><strong>Note:</strong> This model was trained on the SPOC dataset containing C++ code examples.</p>
+    </div>
+    """)
     
     with gr.Row():
+        # Left column - Inputs
         with gr.Column(scale=1):
-            gr.Markdown("### 📝 Input")
-            
             pseudocode_input = gr.Textbox(
-                label="Pseudocode",
-                placeholder="Enter your pseudocode here...\nExample: create integer n",
-                lines=5,
-                max_lines=10
+                label="📝 Pseudocode",
+                placeholder="Enter your pseudocode here... (e.g., 'create string variable s')",
+                lines=3,
+                value="create string s"
             )
             
             with gr.Row():
-                indent_input = gr.Number(
-                    label="Indent Level",
-                    value=1,
-                    precision=0,
-                    info="Indentation level (0=no indent, 1=first level, etc.)"
+                indent_input = gr.Slider(
+                    minimum=1, maximum=10, value=1, step=1,
+                    label="🔢 Indent Level",
+                    info="Indentation level for the code"
                 )
-                
-                line_input = gr.Number(
-                    label="Line Number",
-                    value=1,
-                    precision=0,
+                line_input = gr.Slider(
+                    minimum=1, maximum=100, value=1, step=1,
+                    label="📍 Line Number",
                     info="Line number in the program"
                 )
             
-            gr.Markdown("### ⚙️ Generation Parameters")
+            gr.Markdown("### 🎛️ Generation Parameters")
             
             with gr.Row():
-                temperature_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.5,
-                    value=0.7,
-                    step=0.1,
-                    label="Temperature",
-                    info="Higher = more creative/random"
+                temperature_input = gr.Slider(
+                    minimum=0.1, maximum=2.0, value=0.7, step=0.1,
+                    label="🌡️ Temperature",
+                    info="Higher = more creative, Lower = more focused"
                 )
-                
-                top_p_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.9,
-                    step=0.05,
-                    label="Top-p (Nucleus Sampling)",
-                    info="Probability threshold for sampling"
+                top_p_input = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.9, step=0.05,
+                    label="🎯 Top-p",
+                    info="Nucleus sampling parameter"
                 )
             
-            max_length_slider = gr.Slider(
-                minimum=64,
-                maximum=256,
-                value=128,
-                step=16,
-                label="Max Length",
-                info="Maximum tokens to generate"
+            max_length_input = gr.Slider(
+                minimum=50, maximum=512, value=128, step=10,
+                label="📏 Max Length",
+                info="Maximum number of tokens to generate"
             )
             
             generate_btn = gr.Button("🚀 Generate Code", variant="primary", size="lg")
         
+        # Right column - Output
         with gr.Column(scale=1):
-            gr.Markdown("### 💻 Generated Code")
-            
             output = gr.Textbox(
-                label="Generated Code",
+                label="💻 Generated Code",
                 lines=15,
-                max_lines=20,
+                placeholder="Generated code will appear here...",
                 show_copy_button=True
             )
     
-    gr.Markdown("### 📚 Examples")
-    gr.Examples(
-        examples=examples,
-        inputs=[pseudocode_input, indent_input, line_input, temperature_slider, top_p_slider, max_length_slider],
+    # Examples section
+    gr.Markdown("### 📚 Example Inputs")
+    examples = gr.Examples(
+        examples=create_examples(),
+        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
         outputs=output,
-        fn=gradio_generate_code,
-        cache_examples=False,
+        fn=generate_code,
+        cache_examples=False
     )
     
-    gr.Markdown(
-        """
-        ---
-        ### ℹ️ How to Use:
-        1. **Enter pseudocode**: Write your natural language description
-        2. **Set indent level**: Specify the indentation (0 for no indent, 1 for first level, etc.)
-        3. **Set line number**: Indicate the line position in your program
-        4. **Adjust parameters** (optional): Fine-tune temperature and top-p for different results
-        5. **Click Generate**: Get your code!
-        
-        ### 💡 Tips:
-        - Higher temperature (0.8-1.2) = more creative but potentially less accurate
-        - Lower temperature (0.5-0.7) = more conservative and predictable
-        - Top-p controls diversity; 0.9 is usually a good balance
-        - The model generates C++-style code as it was trained on the SPOC dataset
-        
-        ### 🔗 Resources:
-        - [SPOC Dataset](https://github.com/sumith1896/spoc)
-        - [Research Paper](https://arxiv.org/pdf/1906.04908)
-        - Model trained with LoRA for efficiency
-        """
+    # Event handlers
+    generate_btn.click(
+        fn=generate_code,
+        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
+        outputs=output
     )
     
-    # Connect button to function
-    generate_btn.click(
-        fn=gradio_generate_code,
-        inputs=[pseudocode_input, indent_input, line_input, temperature_slider, top_p_slider, max_length_slider],
+    # Also allow Enter key to generate
+    pseudocode_input.submit(
+        fn=generate_code,
+        inputs=[pseudocode_input, indent_input, line_input, temperature_input, top_p_input, max_length_input],
         outputs=output
     )
+    
+    # Footer
+    gr.HTML("""
+    <div style="text-align: center; margin-top: 30px; padding: 20px; border-top: 1px solid #eee;">
+        <p>🤖 <strong>Model Details:</strong> Fine-tuned GPT-2 with LoRA on SPOC dataset</p>
+        <p>📊 <strong>Training:</strong> Pseudo-code to C++ code generation with structural information</p>
+        <p>⚡ <strong>Powered by:</strong> Transformers, Safetensors, and Gradio</p>
+    </div>
+    """)
 
-# Launch the app
+# Launch configuration
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",  # Required for Hugging Face Spaces
+        server_port=7860,       # Default port for Spaces
+        share=False,            # Don't create public links in Spaces
+        show_api=False,         # Disable API docs for cleaner interface
+        show_error=True,        # Show errors for debugging
+        quiet=False             # Show startup logs
+    )