Create app.py

app.py

@@ -0,0 +1,321 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import numpy as np
+from typing import List, Dict, Tuple
+import json
+import os
+from datetime import datetime
+
+class GRPOTrainer:
+    def __init__(self):
+        self.model = None
+        self.ref_model = None
+        self.tokenizer = None
+        self.optimizer = None
+        self.training_history = []
+        
+    def load_model(self, model_name: str) -> str:
+        """Load the model and tokenizer"""
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            self.model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+            self.ref_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+            
+            # Set padding token
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                
+            # Freeze reference model
+            for param in self.ref_model.parameters():
+                param.requires_grad = False
+                
+            return f"✅ Successfully loaded model: {model_name}"
+        except Exception as e:
+            return f"❌ Error loading model: {str(e)}"
+    
+    def compute_rewards(self, prompts: List[str], responses: List[str]) -> torch.Tensor:
+        """Compute rewards for responses (simplified reward function)"""
+        rewards = []
+        for response in responses:
+            # Simple reward based on response length and diversity
+            length_reward = min(len(response.split()) / 50, 1.0)
+            unique_words = len(set(response.lower().split()))
+            diversity_reward = min(unique_words / 20, 1.0)
+            reward = (length_reward + diversity_reward) / 2
+            rewards.append(reward)
+        return torch.tensor(rewards)
+    
+    def compute_kl_penalty(self, logits: torch.Tensor, ref_logits: torch.Tensor) -> torch.Tensor:
+        """Compute KL divergence penalty"""
+        probs = F.softmax(logits, dim=-1)
+        ref_probs = F.softmax(ref_logits, dim=-1)
+        kl = (probs * (probs / ref_probs).log()).sum(-1)
+        return kl.mean()
+    
+    def grpo_step(self, prompts: List[str], beta: float = 0.1) -> Dict:
+        """Perform one GRPO training step"""
+        if not self.model or not self.tokenizer:
+            return {"error": "Model not loaded"}
+        
+        # Tokenize prompts
+        inputs = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True)
+        
+        # Generate responses
+        with torch.no_grad():
+            outputs = self.model.generate(
+                inputs.input_ids,
+                max_length=inputs.input_ids.shape[1] + 50,
+                do_sample=True,
+                temperature=0.8,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+        
+        # Decode responses
+        responses = []
+        for output in outputs:
+            response = self.tokenizer.decode(output[inputs.input_ids.shape[1]:], skip_special_tokens=True)
+            responses.append(response)
+        
+        # Compute rewards
+        rewards = self.compute_rewards(prompts, responses)
+        
+        # Forward pass through both models
+        self.model.train()
+        model_outputs = self.model(inputs.input_ids)
+        ref_outputs = self.ref_model(inputs.input_ids)
+        
+        # Compute KL penalty
+        kl_penalty = self.compute_kl_penalty(model_outputs.logits, ref_outputs.logits)
+        
+        # Compute loss (simplified GRPO loss)
+        loss = -rewards.mean() + beta * kl_penalty
+        
+        # Backward pass
+        if self.optimizer:
+            self.optimizer.zero_grad()
+            loss.backward()
+            self.optimizer.step()
+        
+        return {
+            "loss": loss.item(),
+            "reward": rewards.mean().item(),
+            "kl_penalty": kl_penalty.item(),
+            "responses": responses
+        }
+    
+    def train(self, prompts: List[str], num_steps: int, lr: float, beta: float) -> str:
+        """Run GRPO training"""
+        if not self.model:
+            return "❌ Please load a model first"
+        
+        # Initialize optimizer
+        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr)
+        
+        results = []
+        for step in range(num_steps):
+            step_result = self.grpo_step(prompts, beta)
+            
+            if "error" in step_result:
+                return f"❌ Error: {step_result['error']}"
+            
+            result_str = f"Step {step + 1}/{num_steps} - Loss: {step_result['loss']:.4f}, Reward: {step_result['reward']:.4f}, KL: {step_result['kl_penalty']:.4f}"
+            results.append(result_str)
+            
+            # Store training history
+            self.training_history.append({
+                "step": step + 1,
+                "loss": step_result['loss'],
+                "reward": step_result['reward'],
+                "kl_penalty": step_result['kl_penalty']
+            })
+        
+        return "\n".join(results)
+    
+    def generate_response(self, prompt: str, max_length: int = 100, temperature: float = 0.8) -> str:
+        """Generate a response using the trained model"""
+        if not self.model or not self.tokenizer:
+            return "❌ Please load a model first"
+        
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                inputs.input_ids,
+                max_length=inputs.input_ids.shape[1] + max_length,
+                temperature=temperature,
+                do_sample=True,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+        
+        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response
+    
+    def save_model(self, save_path: str) -> str:
+        """Save the trained model"""
+        if not self.model:
+            return "❌ No model to save"
+        
+        try:
+            self.model.save_pretrained(save_path)
+            self.tokenizer.save_pretrained(save_path)
+            
+            # Save training history
+            with open(os.path.join(save_path, "training_history.json"), "w") as f:
+                json.dump(self.training_history, f)
+            
+            return f"✅ Model saved to {save_path}"
+        except Exception as e:
+            return f"❌ Error saving model: {str(e)}"
+
+# Initialize trainer
+trainer = GRPOTrainer()
+
+# Gradio interface
+def load_model_interface(model_name):
+    return trainer.load_model(model_name)
+
+def train_interface(prompts_text, num_steps, learning_rate, beta):
+    prompts = [p.strip() for p in prompts_text.split("\n") if p.strip()]
+    if not prompts:
+        return "❌ Please provide at least one prompt"
+    return trainer.train(prompts, int(num_steps), float(learning_rate), float(beta))
+
+def generate_interface(prompt, max_length, temperature):
+    return trainer.generate_response(prompt, int(max_length), float(temperature))
+
+def save_model_interface(save_path):
+    return trainer.save_model(save_path)
+
+def get_training_history():
+    if not trainer.training_history:
+        return "No training history available"
+    
+    history_str = "Training History:\n"
+    history_str += "-" * 50 + "\n"
+    for entry in trainer.training_history[-10:]:  # Show last 10 entries
+        history_str += f"Step {entry['step']}: Loss={entry['loss']:.4f}, Reward={entry['reward']:.4f}, KL={entry['kl_penalty']:.4f}\n"
+    return history_str
+
+# Create Gradio interface
+with gr.Blocks(title="GRPO Model Training") as app:
+    gr.Markdown("# 🚀 GRPO (Group Relative Policy Optimization) Training App")
+    gr.Markdown("Train language models using GRPO technique with this simple interface")
+    
+    with gr.Tab("🔧 Model Setup"):
+        with gr.Row():
+            model_input = gr.Textbox(
+                label="Model Name",
+                value="Palmyra-56b",
+                placeholder="Enter HuggingFace model name (e.g., Palmyra, Qwen, Llama)"
+            )
+            load_btn = gr.Button("Load Model", variant="primary")
+        
+        model_status = gr.Textbox(label="Status", lines=2)
+        load_btn.click(load_model_interface, inputs=model_input, outputs=model_status)
+    
+    with gr.Tab("🎯 Training"):
+        with gr.Row():
+            with gr.Column():
+                prompts_input = gr.Textbox(
+                    label="Training Prompts (one per line)",
+                    lines=5,
+                    value="Tell me about artificial intelligence\nExplain quantum computing\nWhat is machine learning?",
+                    placeholder="Enter your prompts here..."
+                )
+            
+            with gr.Column():
+                num_steps_input = gr.Slider(
+                    label="Number of Training Steps",
+                    minimum=1,
+                    maximum=100,
+                    value=10,
+                    step=1
+                )
+                lr_input = gr.Number(
+                    label="Learning Rate",
+                    value=1e-5,
+                    step=1e-6
+                )
+                beta_input = gr.Number(
+                    label="KL Penalty Weight (β)",
+                    value=0.1,
+                    step=0.01
+                )
+        
+        train_btn = gr.Button("Start Training", variant="primary")
+        training_output = gr.Textbox(label="Training Progress", lines=10)
+        
+        train_btn.click(
+            train_interface,
+            inputs=[prompts_input, num_steps_input, lr_input, beta_input],
+            outputs=training_output
+        )
+    
+    with gr.Tab("💬 Generation"):
+        with gr.Row():
+            with gr.Column():
+                gen_prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    value="Tell me about"
+                )
+                max_length = gr.Slider(
+                    label="Max Length",
+                    minimum=10,
+                    maximum=500,
+                    value=100,
+                    step=10
+                )
+                temp_slider = gr.Slider(
+                    label="Temperature",
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=0.8,
+                    step=0.1
+                )
+            
+            with gr.Column():
+                gen_btn = gr.Button("Generate", variant="primary")
+                gen_output = gr.Textbox(label="Generated Response", lines=10)
+        
+        gen_btn.click(
+            generate_interface,
+            inputs=[gen_prompt, max_length, temp_slider],
+            outputs=gen_output
+        )
+    
+    with gr.Tab("💾 Save Model"):
+        save_path_input = gr.Textbox(
+            label="Save Path",
+            value="./grpo_trained_model",
+            placeholder="Enter path to save the model"
+        )
+        save_btn = gr.Button("Save Model", variant="primary")
+        save_status = gr.Textbox(label="Save Status")
+        
+        save_btn.click(save_model_interface, inputs=save_path_input, outputs=save_status)
+    
+    with gr.Tab("📊 Training History"):
+        history_btn = gr.Button("Refresh History", variant="secondary")
+        history_output = gr.Textbox(label="Training History", lines=15)
+        
+        history_btn.click(get_training_history, outputs=history_output)
+    
+    gr.Markdown("""
+    ## 📝 Instructions:
+    1. **Load Model**: Start by loading a pre-trained model from HuggingFace
+    2. **Training**: Add your prompts and configure training parameters
+    3. **Generation**: Test your trained model with custom prompts
+    4. **Save**: Save your fine-tuned model for later use
+    
+    ## ⚠️ Note:
+    - This is a simplified GRPO implementation for demonstration
+    - For production use, consider more sophisticated reward functions
+    - GPU recommended for larger models
+    """)
+
+# Launch the app
+if __name__ == "__main__":
+    app.launch(share=True)