Update app.py

app.py

@@ -3,357 +3,237 @@ import tensorflow as tf
 import numpy as np
 import os
 import json
-from model import VedaProgrammingLLM, create_veda_model
+
+from model import VedaProgrammingLLM
 from tokenizer import VedaTokenizer
-from train import VedaTrainer
+from train import VedaTrainer, SAMPLE_CODE
 
-# Global variables for model and tokenizer
+# Global state
 model = None
 tokenizer = None
-trainer = None
 
-def load_or_create_model():
-    """Load existing model or create new one"""
-    global model, tokenizer, trainer
+def initialize_model():
+    """Initialize or load model"""
+    global model, tokenizer
     
     model_path = "veda_model"
+    config_file = os.path.join(model_path, "config.json")
     
-    if os.path.exists(os.path.join(model_path, "config.json")):
-        print("Loading existing model...")
-        
-        # Load config
-        with open(os.path.join(model_path, "config.json"), 'r') as f:
-            config = json.load(f)
-        
-        # Load tokenizer
-        tokenizer = VedaTokenizer()
-        tokenizer.load(os.path.join(model_path, "tokenizer.json"))
-        
-        # Create model
-        model = create_veda_model(
-            vocab_size=tokenizer.vocabulary_size,
-            max_length=config.get('max_length', 256),
-            model_size="small"
-        )
-        
-        # Build model
-        dummy_input = tf.zeros((1, config.get('max_length', 256)), dtype=tf.int32)
-        model(dummy_input)
-        
-        # Load weights
-        model.load_weights(os.path.join(model_path, "model_weights.h5"))
-        print("Model loaded successfully!")
-    else:
-        print("No existing model found. Creating and training new model...")
-        trainer = VedaTrainer(
-            data_path="programming.txt",
-            vocab_size=5000,
-            max_length=256,
-            batch_size=8,
-            model_size="small"
-        )
-        trainer.train(epochs=5, save_path=model_path)
+    try:
+        if os.path.exists(config_file):
+            print("Loading saved model...")
+            
+            with open(config_file, 'r') as f:
+                config = json.load(f)
+            
+            tokenizer = VedaTokenizer()
+            tokenizer.load(os.path.join(model_path, "tokenizer.json"))
+            
+            model = VedaProgrammingLLM(
+                vocab_size=config['vocab_size'],
+                max_length=config['max_length'],
+                d_model=config['d_model'],
+                num_heads=config['num_heads'],
+                num_layers=config['num_layers'],
+                ff_dim=config['ff_dim']
+            )
+            
+            dummy = tf.zeros((1, config['max_length']), dtype=tf.int32)
+            model(dummy)
+            model.load_weights(os.path.join(model_path, "weights.h5"))
+            print("Model loaded!")
+        else:
+            print("Training new model...")
+            trainer = VedaTrainer(
+                data_path="programming.txt",
+                vocab_size=3000,
+                max_length=128,
+                batch_size=8
+            )
+            trainer.train(epochs=5, save_path=model_path)
+            model = trainer.model
+            tokenizer = trainer.tokenizer
+            print("Model trained!")
+            
+    except Exception as e:
+        print(f"Error: {e}")
+        print("Creating fresh model...")
+        trainer = VedaTrainer()
+        trainer.train(epochs=5)
         model = trainer.model
         tokenizer = trainer.tokenizer
-        print("Model trained and ready!")
 
-def generate_code(
-    prompt: str,
-    max_tokens: int,
-    temperature: float,
-    top_k: int,
-    top_p: float
-) -> str:
-    """Generate code based on prompt"""
+def generate_code(prompt: str, max_tokens: int, temperature: float, top_k: int) -> str:
+    """Generate code from prompt"""
     global model, tokenizer
     
     if model is None or tokenizer is None:
-        return "Error: Model not loaded. Please wait for initialization."
+        return "Model not loaded. Please wait..."
     
     try:
-        # Encode prompt
-        prompt_tokens = tokenizer.encode(prompt)
+        if not prompt.strip():
+            return "Please enter a prompt."
         
-        if len(prompt_tokens) == 0:
-            return "Error: Could not tokenize prompt."
+        tokens = tokenizer.encode(prompt)
+        if len(tokens) == 0:
+            tokens = [2]  # START token
         
-        # Generate
-        generated_tokens = model.generate(
-            prompt_tokens,
+        generated = model.generate(
+            tokens,
             max_new_tokens=int(max_tokens),
             temperature=float(temperature),
-            top_k=int(top_k),
-            top_p=float(top_p)
+            top_k=int(top_k)
         )
         
-        # Decode
-        generated_text = tokenizer.decode(generated_tokens)
+        result = tokenizer.decode(generated)
+        return result
         
-        return generated_text
-    
     except Exception as e:
-        return f"Error generating code: {str(e)}"
+        return f"Error: {str(e)}"
 
-def train_model(
-    training_data: str,
-    epochs: int,
-    learning_rate: float
-) -> str:
-    """Train or fine-tune the model"""
-    global model, tokenizer, trainer
+def train_on_data(training_data: str, epochs: int) -> str:
+    """Train model on provided data"""
+    global model, tokenizer
     
     try:
-        # Save training data
-        with open("programming.txt", 'w', encoding='utf-8') as f:
+        with open("programming.txt", 'w') as f:
             f.write(training_data)
         
-        # Create trainer
         trainer = VedaTrainer(
             data_path="programming.txt",
-            vocab_size=5000,
-            max_length=256,
-            batch_size=8,
-            model_size="small"
+            vocab_size=3000,
+            max_length=128,
+            batch_size=8
         )
         
-        # Train
         history = trainer.train(epochs=int(epochs), save_path="veda_model")
         
-        # Update global references
         model = trainer.model
         tokenizer = trainer.tokenizer
         
-        # Get final metrics
         final_loss = history.history['loss'][-1]
         final_acc = history.history.get('accuracy', [0])[-1]
         
-        return f"""Training completed successfully!
-        
-Final Loss: {final_loss:.4f}
-Final Accuracy: {final_acc:.4f}
+        return f"""✅ Training Complete!
+
+Loss: {final_loss:.4f}
+Accuracy: {final_acc:.4f}
 Epochs: {epochs}
-Vocabulary Size: {tokenizer.vocabulary_size}
+Vocab Size: {tokenizer.vocabulary_size}
 """
-    
     except Exception as e:
-        return f"Error during training: {str(e)}"
+        return f"❌ Training Error: {str(e)}"
 
 def get_model_info() -> str:
-    """Get information about the current model"""
+    """Get model information"""
     global model, tokenizer
     
     if model is None:
         return "No model loaded."
     
     config = model.get_config()
+    params = model.count_params()
     
-    info = f"""
-**Veda Programming LLM**
-
-- Vocabulary Size: {config['vocab_size']}
-- Max Sequence Length: {config['max_length']}
-- Model Dimension: {config['d_model']}
-- Number of Heads: {config['num_heads']}
-- Number of Layers: {config['num_layers']}
-- Feed Forward Dimension: {config['ff_dim']}
+    return f"""## 🕉️ Veda Programming LLM
 
-Total Parameters: {model.count_params():,}
+| Property | Value |
+|----------|-------|
+| Vocabulary Size | {config['vocab_size']} |
+| Max Length | {config['max_length']} |
+| Model Dimension | {config['d_model']} |
+| Attention Heads | {config['num_heads']} |
+| Transformer Layers | {config['num_layers']} |
+| FFN Dimension | {config['ff_dim']} |
+| **Total Parameters** | **{params:,}** |
 """
-    return info
 
-# Create Gradio interface
-def create_interface():
-    """Create the Gradio interface"""
-    
-    with gr.Blocks(
-        title="Veda Programming LLM",
-        theme=gr.themes.Soft()
-    ) as demo:
+# Build interface
+def create_app():
+    with gr.Blocks(title="Veda Programming", theme=gr.themes.Soft()) as app:
         
         gr.Markdown("""
         # 🕉️ Veda Programming LLM
-        
-        A TensorFlow-based language model for code generation.
-        
-        Train on your own code or use the pre-trained model to generate programming code.
+        ### TensorFlow-based Code Generation Model
         """)
         
         with gr.Tabs():
             # Generation Tab
-            with gr.TabItem("💻 Code Generation"):
+            with gr.TabItem("💻 Generate Code"):
                 with gr.Row():
-                    with gr.Column(scale=1):
-                        prompt_input = gr.Textbox(
+                    with gr.Column():
+                        prompt = gr.Textbox(
                             label="Code Prompt",
                             placeholder="def fibonacci(",
-                            lines=5,
+                            lines=3,
                             value="def calculate_sum("
                         )
                         
                         with gr.Row():
-                            max_tokens = gr.Slider(
-                                minimum=10,
-                                maximum=500,
-                                value=100,
-                                step=10,
-                                label="Max New Tokens"
-                            )
-                            temperature = gr.Slider(
-                                minimum=0.1,
-                                maximum=2.0,
-                                value=0.7,
-                                step=0.1,
-                                label="Temperature"
-                            )
+                            max_tokens = gr.Slider(10, 200, value=50, step=5, label="Max Tokens")
+                            temperature = gr.Slider(0.1, 1.5, value=0.8, step=0.1, label="Temperature")
                         
-                        with gr.Row():
-                            top_k = gr.Slider(
-                                minimum=1,
-                                maximum=100,
-                                value=50,
-                                step=1,
-                                label="Top-K"
-                            )
-                            top_p = gr.Slider(
-                                minimum=0.1,
-                                maximum=1.0,
-                                value=0.9,
-                                step=0.05,
-                                label="Top-P (Nucleus)"
-                            )
+                        top_k = gr.Slider(1, 100, value=40, step=5, label="Top-K")
                         
-                        generate_btn = gr.Button(
-                            "🚀 Generate Code",
-                            variant="primary"
-                        )
+                        gen_btn = gr.Button("🚀 Generate", variant="primary")
                     
-                    with gr.Column(scale=1):
-                        output_code = gr.Code(
-                            label="Generated Code",
-                            language="python",
-                            lines=15
-                        )
+                    with gr.Column():
+                        output = gr.Code(label="Generated Code", language="python", lines=12)
+                
+                gen_btn.click(
+                    generate_code,
+                    inputs=[prompt, max_tokens, temperature, top_k],
+                    outputs=output
+                )
                 
-                generate_btn.click(
-                    fn=generate_code,
-                    inputs=[prompt_input, max_tokens, temperature, top_k, top_p],
-                    outputs=output_code
+                gr.Examples(
+                    examples=[
+                        ["def fibonacci(n):", 60, 0.7, 40],
+                        ["def bubble_sort(arr):", 80, 0.8, 40],
+                        ["class Calculator:", 100, 0.7, 50],
+                        ["def binary_search(", 70, 0.8, 40],
+                    ],
+                    inputs=[prompt, max_tokens, temperature, top_k],
+                    outputs=output,
+                    fn=generate_code
                 )
             
             # Training Tab
-            with gr.TabItem("🎓 Training"):
-                gr.Markdown("""
-                ### Train the Model
-                
-                Paste your programming code below to train or fine-tune the model.
-                """)
-                
-                training_data = gr.Textbox(
-                    label="Training Data (Python Code)",
-                    placeholder="Paste your Python code here...",
-                    lines=15,
-                    value='''def hello_world():
-    print("Hello, World!")
-    return True
-
-def add_numbers(a, b):
-    return a + b
-
-class MyClass:
-    def __init__(self, value):
-        self.value = value
-    
-    def get_value(self):
-        return self.value
-'''
+            with gr.TabItem("🎓 Train Model"):
+                training_input = gr.Textbox(
+                    label="Training Code",
+                    placeholder="Paste Python code...",
+                    lines=12,
+                    value=SAMPLE_CODE[:1500]
                 )
                 
-                with gr.Row():
-                    epochs_input = gr.Slider(
-                        minimum=1,
-                        maximum=50,
-                        value=10,
-                        step=1,
-                        label="Epochs"
-                    )
-                    lr_input = gr.Slider(
-                        minimum=0.00001,
-                        maximum=0.01,
-                        value=0.0001,
-                        step=0.00001,
-                        label="Learning Rate"
-                    )
+                epochs_slider = gr.Slider(1, 20, value=5, step=1, label="Epochs")
                 
-                train_btn = gr.Button(
-                    "🎯 Start Training",
-                    variant="primary"
-                )
-                
-                training_output = gr.Textbox(
-                    label="Training Output",
-                    lines=10
-                )
+                train_btn = gr.Button("🎯 Train Model", variant="primary")
+                train_output = gr.Textbox(label="Training Results", lines=8)
                 
                 train_btn.click(
-                    fn=train_model,
-                    inputs=[training_data, epochs_input, lr_input],
-                    outputs=training_output
+                    train_on_data,
+                    inputs=[training_input, epochs_slider],
+                    outputs=train_output
                 )
             
-            # Model Info Tab
+            # Info Tab
             with gr.TabItem("ℹ️ Model Info"):
-                model_info = gr.Markdown()
-                refresh_btn = gr.Button("🔄 Refresh Info")
-                refresh_btn.click(
-                    fn=get_model_info,
-                    outputs=model_info
-                )
-            
-            # Examples Tab
-            with gr.TabItem("📝 Examples"):
-                gr.Markdown("""
-                ### Example Prompts
-                
-                Try these prompts for code generation:
-                
-                1. `def fibonacci(`
-                2. `class Calculator:`
-                3. `def binary_search(arr, target):`
-                4. `async def fetch_data(`
-                5. `def quicksort(arr):`
-                6. `class LinkedList:`
-                """)
-                
-                gr.Examples(
-                    examples=[
-                        ["def fibonacci(", 100, 0.7, 50, 0.9],
-                        ["class Calculator:", 150, 0.8, 40, 0.95],
-                        ["def binary_search(arr, target):", 100, 0.6, 50, 0.9],
-                        ["def bubble_sort(", 100, 0.7, 50, 0.9],
-                    ],
-                    inputs=[prompt_input, max_tokens, temperature, top_k, top_p],
-                    outputs=output_code,
-                    fn=generate_code
-                )
+                info_output = gr.Markdown()
+                refresh_btn = gr.Button("🔄 Refresh")
+                refresh_btn.click(get_model_info, outputs=info_output)
         
         gr.Markdown("""
         ---
-        **Veda Programming LLM** - Built with TensorFlow & Gradio
-        
-        *Note: This is a demonstration model. For production use, train on larger datasets.*
+        **Veda Programming LLM** | Built with TensorFlow & Gradio
         """)
     
-    return demo
+    return app
 
-# Main execution
-if __name__ == "__main__":
-    print("Initializing Veda Programming LLM...")
-    load_or_create_model()
-    
-    print("Starting Gradio interface...")
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+# Main
+print("🕉️ Initializing Veda Programming LLM...")
+initialize_model()
+
+print("🚀 Starting Gradio...")
+app = create_app()
+app.launch(server_name="0.0.0.0", server_port=7860)