Upload 12 files

app.py

@@ -0,0 +1,227 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torch
+from rich.console import Console
+import time
+
+# Initialize rich console for better logging
+console = Console()
+
+# Load the model and tokenizer with the same configuration as training
+console.print("[bold green]Loading model and tokenizer...[/bold green]")
+
+# Configure 4-bit quantization with memory optimizations
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_quant_storage=torch.float16,
+)
+
+# Load model with quantization and memory optimizations
+model = AutoModelForCausalLM.from_pretrained(
+    "./fine-tuned-model",
+    quantization_config=quantization_config,
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+)
+tokenizer = AutoTokenizer.from_pretrained("./fine-tuned-model")
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = 'left'
+
+def generate_response(
+    prompt,
+    max_length=128,  # Match training max_length
+    temperature=0.7,
+    top_p=0.9,
+    num_generations=2,  # Match training num_generations
+    repetition_penalty=1.1,
+    do_sample=True,
+):
+    try:
+        # Get the device of the model
+        device = next(model.parameters()).device
+        
+        # Tokenize the input
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+        
+        # Move inputs to the same device as the model
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        
+        # Generate response
+        with torch.no_grad():  # Disable gradient computation
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_length,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_p=top_p,
+                num_return_sequences=num_generations,
+                repetition_penalty=repetition_penalty,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        
+        # Decode and return the responses
+        responses = []
+        for output in outputs:
+            response = tokenizer.decode(output, skip_special_tokens=True)
+            responses.append(response)
+        
+        return "\n\n---\n\n".join(responses)
+    except Exception as e:
+        console.print(f"[bold red]Error during generation: {str(e)}[/bold red]")
+        return f"Error: {str(e)}"
+
+# Create custom CSS for better UI
+custom_css = """
+.gradio-container {
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+}
+.container {
+    max-width: 800px;
+    margin: auto;
+    padding: 20px;
+}
+.title {
+    text-align: center;
+    color: #2c3e50;
+    margin-bottom: 20px;
+}
+.description {
+    color: #34495e;
+    line-height: 1.6;
+    margin-bottom: 20px;
+}
+"""
+
+# Create the Gradio interface with enhanced UI
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # Phi-2 Fine-tuned with GRPO and qLoRA
+        This model has been fine-tuned using GRPO (Generative Reward-Penalized Optimization) and compressed using qLoRA.
+        Try it out with different prompts and generation parameters!
+        """,
+        elem_classes="title"
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Enter your prompt here...",
+                lines=3,
+                show_label=True,
+            )
+            
+            with gr.Row():
+                with gr.Column():
+                    max_length = gr.Slider(
+                        minimum=32,
+                        maximum=256,
+                        value=128,
+                        step=32,
+                        label="Max Length",
+                        info="Maximum number of tokens to generate"
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.7,
+                        step=0.1,
+                        label="Temperature",
+                        info="Higher values make output more random, lower values more deterministic"
+                    )
+                with gr.Column():
+                    top_p = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.9,
+                        step=0.1,
+                        label="Top-p",
+                        info="Nucleus sampling parameter"
+                    )
+                    num_generations = gr.Slider(
+                        minimum=1,
+                        maximum=4,
+                        value=2,
+                        step=1,
+                        label="Number of Generations",
+                        info="Number of different responses to generate"
+                    )
+            
+            with gr.Row():
+                with gr.Column():
+                    repetition_penalty = gr.Slider(
+                        minimum=1.0,
+                        maximum=2.0,
+                        value=1.1,
+                        step=0.1,
+                        label="Repetition Penalty",
+                        info="Higher values prevent repetition"
+                    )
+                with gr.Column():
+                    do_sample = gr.Checkbox(
+                        value=True,
+                        label="Enable Sampling",
+                        info="Enable/disable sampling for deterministic output"
+                    )
+            
+            generate_btn = gr.Button("Generate", variant="primary")
+        
+        with gr.Column(scale=3):
+            output = gr.Textbox(
+                label="Generated Response(s)",
+                lines=10,
+                show_label=True,
+            )
+    
+    gr.Markdown(
+        """
+        ### Example Prompts
+        Try these example prompts to test the model:
+        
+        1. **Technical Question**: "What is machine learning?"
+        2. **Creative Writing**: "Write a short story about a robot learning to paint."
+        3. **Technical Explanation**: "Explain quantum computing in simple terms."
+        4. **Creative Writing**: "Write a poem about artificial intelligence."
+        """,
+        elem_classes="description"
+    )
+    
+    # Add examples
+    gr.Examples(
+        examples=[
+            ["What is machine learning?"],
+            ["Write a short story about a robot learning to paint."],
+            ["Explain quantum computing in simple terms."],
+            ["Write a poem about artificial intelligence."]
+        ],
+        inputs=prompt
+    )
+    
+    # Connect the interface
+    generate_btn.click(
+        fn=generate_response,
+        inputs=[
+            prompt,
+            max_length,
+            temperature,
+            top_p,
+            num_generations,
+            repetition_penalty,
+            do_sample
+        ],
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    console.print("[bold green]Starting Gradio interface...[/bold green]")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=True  # Enable sharing for HuggingFace Spaces
+    ) 

fine-tuned-model/README.md

@@ -0,0 +1,202 @@
+---
+base_model: microsoft/phi-2
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.14.0

fine-tuned-model/adapter_config.json

@@ -0,0 +1,34 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "microsoft/phi-2",
+  "bias": "none",
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "q_proj",
+    "v_proj",
+    "k_proj",
+    "o_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

fine-tuned-model/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1491777fce601097323585d69288bc5376d647628f0673c54196edaf47692f9
+size 31483040

fine-tuned-model/added_tokens.json

@@ -0,0 +1,40 @@
+{
+  "\t\t": 50294,
+  "\t\t\t": 50293,
+  "\t\t\t\t": 50292,
+  "\t\t\t\t\t": 50291,
+  "\t\t\t\t\t\t": 50290,
+  "\t\t\t\t\t\t\t": 50289,
+  "\t\t\t\t\t\t\t\t": 50288,
+  "\t\t\t\t\t\t\t\t\t": 50287,
+  "  ": 50286,
+  "   ": 50285,
+  "    ": 50284,
+  "     ": 50283,
+  "      ": 50282,
+  "       ": 50281,
+  "        ": 50280,
+  "         ": 50279,
+  "          ": 50278,
+  "           ": 50277,
+  "            ": 50276,
+  "             ": 50275,
+  "              ": 50274,
+  "               ": 50273,
+  "                ": 50272,
+  "                 ": 50271,
+  "                  ": 50270,
+  "                   ": 50269,
+  "                    ": 50268,
+  "                     ": 50267,
+  "                      ": 50266,
+  "                       ": 50265,
+  "                        ": 50264,
+  "                         ": 50263,
+  "                          ": 50262,
+  "                           ": 50261,
+  "                            ": 50260,
+  "                             ": 50259,
+  "                              ": 50258,
+  "                               ": 50257
+}

fine-tuned-model/special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>",
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

fine-tuned-model/tokenizer_config.json

@@ -0,0 +1,326 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "50257": {
+      "content": "                               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50258": {
+      "content": "                              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50259": {
+      "content": "                             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50260": {
+      "content": "                            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50261": {
+      "content": "                           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50262": {
+      "content": "                          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50263": {
+      "content": "                         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50264": {
+      "content": "                        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50265": {
+      "content": "                       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50266": {
+      "content": "                      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50267": {
+      "content": "                     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50268": {
+      "content": "                    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50269": {
+      "content": "                   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50270": {
+      "content": "                  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50271": {
+      "content": "                 ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50272": {
+      "content": "                ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50273": {
+      "content": "               ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50274": {
+      "content": "              ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50275": {
+      "content": "             ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50276": {
+      "content": "            ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50277": {
+      "content": "           ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50278": {
+      "content": "          ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50279": {
+      "content": "         ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50280": {
+      "content": "        ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50281": {
+      "content": "       ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50282": {
+      "content": "      ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50283": {
+      "content": "     ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50284": {
+      "content": "    ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50285": {
+      "content": "   ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50286": {
+      "content": "  ",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50287": {
+      "content": "\t\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50288": {
+      "content": "\t\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50289": {
+      "content": "\t\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50290": {
+      "content": "\t\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50291": {
+      "content": "\t\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50292": {
+      "content": "\t\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50293": {
+      "content": "\t\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "50294": {
+      "content": "\t\t",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 2048,
+  "pad_token": "<|endoftext|>",
+  "return_token_type_ids": false,
+  "tokenizer_class": "CodeGenTokenizer",
+  "unk_token": "<|endoftext|>"
+}

requirements.txt

@@ -0,0 +1,12 @@
+torch>=2.0.0
+transformers>=4.30.0
+datasets>=2.12.0
+accelerate>=0.20.0
+bitsandbytes>=0.41.0
+peft>=0.4.0
+pytorch-lightning>=2.0.0
+gradio>=3.40.0
+wandb>=0.15.0
+rich>=13.0.0
+sentencepiece>=0.1.99
+protobuf>=4.23.0 

train.py

@@ -0,0 +1,380 @@
+import os
+import torch
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.loggers import WandbLogger
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    get_linear_schedule_with_warmup,
+    BitsAndBytesConfig,
+    TrainingArguments,
+)
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from rich.console import Console
+from torch.utils.data import Dataset, DataLoader
+
+# Enable Tensor Core optimization for RTX GPUs
+torch.set_float32_matmul_precision('medium')
+
+# Initialize rich console for better logging
+console = Console()
+
+class TextDataset(Dataset):
+    def __init__(self, dataset, tokenizer, max_length=512):
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        # Ensure tokenizer has a padding token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        # Combine instruction and input if they exist
+        prompt = item.get("instruction", "")
+        if item.get("input"):
+            prompt += "\n" + item["input"]
+        
+        # Tokenize the prompt
+        encoding = self.tokenizer(
+            prompt,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        return {
+            "input_ids": encoding["input_ids"].squeeze(),
+            "attention_mask": encoding["attention_mask"].squeeze(),
+            "prompt": prompt
+        }
+
+class GRPOModel(pl.LightningModule):
+    def __init__(
+        self,
+        model_name="microsoft/phi-2",
+        learning_rate=2e-5,
+        num_train_epochs=3,
+        warmup_steps=100,
+        batch_size=2,
+        max_length=128,
+        beta=0.04,
+        num_generations=2,
+        train_dataset=None,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+        
+        # Store train dataset
+        self.train_dataset = train_dataset
+        
+        # Configure 4-bit quantization with memory optimizations
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_quant_storage=torch.float16,
+        )
+        
+        # Load model with quantization and memory optimizations
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=quantization_config,
+            device_map="auto",
+            trust_remote_code=True,
+            torch_dtype=torch.float16,
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = 'left'
+        
+        # Prepare model for training
+        self.model = prepare_model_for_kbit_training(self.model)
+        
+        # LoRA configuration
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+        
+        self.model = get_peft_model(self.model, lora_config)
+        
+        # Store model name for reference model
+        self.model_name = model_name
+        self.ref_model = None
+
+    def setup(self, stage=None):
+        # Move model to the correct device after initialization
+        if stage == "fit":
+            self.model = self.model.to(self.device)
+
+    def get_reference_model(self):
+        if self.ref_model is None:
+            # Load reference model with quantization
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+            )
+            self.ref_model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                quantization_config=quantization_config,
+                device_map=None,
+                trust_remote_code=True,
+            )
+            self.ref_model.eval()
+            self.ref_model = self.ref_model.to(self.device)
+        return self.ref_model
+
+    def reward_function(self, completions):
+        rewards = []
+        for completion in completions:
+            # Reward based on length (normalized)
+            length_reward = len(completion.split()) / 100
+            
+            # Reward based on diversity (unique words)
+            unique_words = len(set(completion.lower().split()))
+            diversity_reward = unique_words / len(completion.split())
+            
+            # Combined reward
+            reward = 0.7 * length_reward + 0.3 * diversity_reward
+            rewards.append(reward)
+        
+        return torch.tensor(rewards, device=self.device)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True
+        )
+        return outputs.logits
+
+    def training_step(self, batch, batch_idx):
+        # Generate completions
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        prompts = batch["prompt"]
+        
+        # Generate multiple completions for each prompt
+        all_completions = []
+        for _ in range(self.hparams.num_generations):
+            outputs = self.model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=128,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                pad_token_id=self.tokenizer.eos_token_id
+            )
+            completions = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+            all_completions.extend(completions)
+        
+        # Calculate rewards
+        rewards = self.reward_function(all_completions)
+        
+        # Calculate KL divergence
+        ref_model = self.get_reference_model()
+        with torch.no_grad():
+            ref_outputs = ref_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                return_dict=True
+            )
+            ref_logits = ref_outputs.logits
+        
+        policy_logits = self(input_ids, attention_mask)
+        kl_div = torch.nn.functional.kl_div(
+            torch.nn.functional.log_softmax(policy_logits, dim=-1),
+            torch.nn.functional.softmax(ref_logits, dim=-1),
+            reduction='batchmean'
+        )
+        
+        # Calculate GRPO loss
+        loss = -rewards.mean() + self.hparams.beta * kl_div
+        
+        self.log("train_loss", loss)
+        self.log("train_reward", rewards.mean())
+        self.log("train_kl_div", kl_div)
+        
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.learning_rate)
+        scheduler = get_linear_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=self.hparams.warmup_steps,
+            num_training_steps=self.hparams.num_train_epochs * len(self.train_dataloader())
+        )
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": scheduler,
+                "monitor": "train_loss",
+                "frequency": 1
+            }
+        }
+
+    def on_train_end(self):
+        # Clean up reference model to free memory
+        if self.ref_model is not None:
+            del self.ref_model
+            self.ref_model = None
+            torch.cuda.empty_cache()
+
+    def train_dataloader(self):
+        if self.train_dataset is None:
+            raise ValueError("Train dataset not provided")
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.hparams.batch_size,
+            shuffle=True,
+            num_workers=4,
+            persistent_workers=True,
+            pin_memory=True
+        )
+
+class TextDataModule(pl.LightningDataModule):
+    def __init__(
+        self,
+        tokenizer,
+        max_length=256,
+        batch_size=4,
+        num_workers=4,
+        pin_memory=True,
+    ):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.pin_memory = pin_memory
+
+def main():
+    # Load dataset
+    dataset = load_dataset("tatsu-lab/alpaca")
+    train_dataset = dataset["train"].select(range(500))
+    
+    # Initialize tokenizer with left padding
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = 'left'
+    
+    # Create dataset with reduced max length
+    train_dataset = TextDataset(train_dataset, tokenizer, max_length=128)
+    
+    # Initialize model with optimized parameters for RTX 4060 Laptop
+    model = GRPOModel(
+        train_dataset=train_dataset,
+        batch_size=2,
+        num_generations=2,
+        max_length=128,
+        learning_rate=1e-5,
+        beta=0.02,
+    )
+    
+    # Initialize logger and callbacks
+    wandb_logger = WandbLogger(project="llm-finetuning")
+    checkpoint_callback = ModelCheckpoint(
+        dirpath="./checkpoints",
+        filename="model-{epoch:02d}-{step:04d}",
+        monitor="train_loss",
+        mode="min",
+        save_top_k=3,
+    )
+    early_stopping = EarlyStopping(
+        monitor="train_loss",
+        patience=3,
+        mode="min",
+    )
+    
+    # Training configuration
+    training_args = TrainingArguments(
+        output_dir="./fine-tuned-model",
+        num_train_epochs=3,
+        per_device_train_batch_size=2,
+        gradient_accumulation_steps=4,
+        learning_rate=1e-5,
+        weight_decay=0.01,
+        warmup_steps=50,
+        logging_steps=10,
+        save_strategy="epoch",
+        evaluation_strategy="no",
+        fp16=False,
+        gradient_checkpointing=True,
+        optim="adamw_torch",
+        lr_scheduler_type="cosine",
+        remove_unused_columns=False,
+        report_to="wandb",
+        dataloader_num_workers=4,
+        dataloader_pin_memory=True,
+        torch_compile=True,
+        max_grad_norm=1.0,
+        group_by_length=True,
+    )
+    
+    # Initialize trainer with memory-optimized settings
+    trainer = pl.Trainer(
+        max_epochs=3,
+        accelerator="gpu",
+        devices=1,
+        precision="32",
+        gradient_clip_val=1.0,
+        accumulate_grad_batches=4,
+        log_every_n_steps=10,
+        val_check_interval=0.5,
+        callbacks=[
+            checkpoint_callback,
+            early_stopping,
+        ],
+        strategy="auto",
+    )
+    
+    # Train the model
+    console.print("[bold green]Starting training...[/bold green]")
+    console.print("[bold yellow]Training with optimized settings for RTX 4060 Laptop GPU[/bold yellow]")
+    console.print(f"Batch size: {model.hparams.batch_size}")
+    console.print(f"Generations per prompt: {model.hparams.num_generations}")
+    console.print(f"Max sequence length: {model.hparams.max_length}")
+    trainer.fit(model)
+    console.print("[bold green]Training completed![/bold green]")
+    
+    # Save the model
+    model.model.save_pretrained("./fine-tuned-model")
+    model.tokenizer.save_pretrained("./fine-tuned-model")
+    console.print("[bold green]Model saved successfully![/bold green]")
+    
+    # Test the model
+    test_prompt = "What is machine learning?"
+    console.print("\n[bold blue]Testing the model:[/bold blue]")
+    console.print(f"Original prompt: {test_prompt}")
+    
+    inputs = model.tokenizer(test_prompt, return_tensors="pt").to(model.device)
+    outputs = model.model.generate(
+        **inputs,
+        max_new_tokens=128,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+    )
+    response = model.tokenizer.decode(outputs[0], skip_special_tokens=True)
+    console.print(f"Generated response: {response}")
+
+if __name__ == "__main__":
+    main()