{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.14"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84795,"databundleVersionId":10462807,"sourceType":"competition"},{"sourceId":166218,"sourceType":"modelInstanceVersion","modelInstanceId":141432,"modelId":164048}],"dockerImageVersionId":30805,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":22.341371,"end_time":"2024-12-11T03:22:13.479076","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-12-11T03:21:51.137705","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import pandas as pd\nfrom typing import Dict, List\nimport zipfile\nimport re\nimport os\n\ndef unzip_data():\n    \"\"\"Unzip the data.a_zip file\"\"\"\n    try:\n        with zipfile.ZipFile('/kaggle/input/konwinski-prize/data.a_zip', 'r') as zip_ref:\n            zip_ref.extractall('/kaggle/working/data')\n        print(\"Data successfully unzipped!\")\n    except Exception as e:\n        print(f\"Error unzipping data: {e}\")\n\nclass DataProcessor:\n    def __init__(self, df: pd.DataFrame):\n        self.df = df\n        \n    def extract_code_snippets(self, problem_statement: str) -> List[str]:\n        \"\"\"Extract code snippets between triple backticks\"\"\"\n        try:\n            code_pattern = r\"```(?:py|python)?\\n(.*?)\\n```\"\n            return re.findall(code_pattern, str(problem_statement), re.DOTALL)\n        except Exception as e:\n            print(f\"Error extracting code snippets: {e}\")\n            return []\n    \n    def parse_patch(self, patch: str) -> Dict:\n        \"\"\"Parse git diff to understand file changes\"\"\"\n        try:\n            files_changed = []\n            for line in str(patch).split('\\n'):\n                if line.startswith('diff --git'):\n                    files_changed.append(line.split()[-1][2:])\n            return {\n                'files_changed': files_changed,\n                'raw_patch': patch\n            }\n        except Exception as e:\n            print(f\"Error parsing patch: {e}\")\n            return {'files_changed': [], 'raw_patch': ''}\n    \n    def process_single_issue(self, row) -> Dict:\n        \"\"\"Process a single GitHub issue\"\"\"\n        try:\n            try:\n                pass_to_pass = eval(str(row['PASS_TO_PASS']).replace('\\x00', ''))\n            except:\n                pass_to_pass = []\n                \n            try:\n                fail_to_pass = eval(str(row['FAIL_TO_PASS']).replace('\\x00', ''))\n            except:\n                fail_to_pass = []\n                \n            return {\n                'id': str(row['instance_id']),\n                'repo': str(row['repo']),\n                'problem': {\n                    'description': str(row['problem_statement']),\n                    'code_snippets': self.extract_code_snippets(str(row['problem_statement'])),\n                    'error_type': self.extract_error_type(str(row['problem_statement']))\n                },\n                'solution': self.parse_patch(str(row['patch'])),\n                'tests': {\n                    'pass_to_pass': pass_to_pass,\n                    'fail_to_pass': fail_to_pass\n                }\n            }\n        except Exception as e:\n            print(f\"Error processing single issue: {e}\")\n            return {}\n    \n    def extract_error_type(self, problem_statement: str) -> str:\n        \"\"\"Extract the type of error from problem statement\"\"\"\n        try:\n            error_pattern = r\"([A-Za-z]+Error:)\"\n            matches = re.findall(error_pattern, str(problem_statement))\n            return matches[0] if matches else \"Unknown\"\n        except Exception as e:\n            print(f\"Error extracting error type: {e}\")\n            return \"Unknown\"\n    \n    def process_all_data(self) -> List[Dict]:\n        \"\"\"Process all issues in the dataset\"\"\"\n        processed_data = []\n        for _, row in self.df.iterrows():\n            try:\n                processed_issue = self.process_single_issue(row)\n                if processed_issue:  # Only add if not empty\n                    processed_data.append(processed_issue)\n            except Exception as e:\n                print(f\"Error processing row: {e}\")\n                continue\n        return processed_data\n\ndef load_and_process_data():\n    \"\"\"Load and process the data\"\"\"\n    try:\n        # First unzip the data\n        unzip_data()\n        \n        # Read the parquet file with handling for null bytes\n        df = pd.read_parquet('/kaggle/working/data/data/data.parquet')\n        \n        # Clean null bytes from string columns\n        string_columns = df.select_dtypes(include=['object']).columns\n        for col in string_columns:\n            df[col] = df[col].astype(str).str.replace('\\x00', '')\n        \n        # Create processor instance and process data\n        processor = DataProcessor(df)\n        processed_data = processor.process_all_data()\n        \n        if processed_data:\n            print(f\"\\nTotal issues processed: {len(processed_data)}\")\n            error_types = [issue['problem']['error_type'] for issue in processed_data]\n            print(f\"Error types found: {set(error_types)}\")\n        \n        return processed_data\n    \n    except Exception as e:\n        print(f\"Error in data processing: {e}\")\n        return None\n\nif __name__ == \"__main__\":\n    processed_data = load_and_process_data()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T15:39:17.232745Z","iopub.execute_input":"2025-01-09T15:39:17.233095Z","iopub.status.idle":"2025-01-09T15:39:21.388976Z","shell.execute_reply.started":"2025-01-09T15:39:17.233043Z","shell.execute_reply":"2025-01-09T15:39:21.387794Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install -U bitsandbytes\n!pip install peft\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T15:39:21.390606Z","iopub.execute_input":"2025-01-09T15:39:21.391006Z","iopub.status.idle":"2025-01-09T15:39:41.435202Z","shell.execute_reply.started":"2025-01-09T15:39:21.390977Z","shell.execute_reply":"2025-01-09T15:39:41.43378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom transformers import (\n    AutoConfig,\n    AutoTokenizer,\n    AutoModelForCausalLM,\n    BitsAndBytesConfig,\n    TrainingArguments,\n    Trainer,\n    DataCollatorForLanguageModeling,\n)\nfrom peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training\nfrom datasets import Dataset\nfrom tqdm import tqdm\nimport gc\nfrom sklearn.model_selection import train_test_split\nimport logging\nimport warnings\nfrom difflib import SequenceMatcher\nimport numpy as np\n\nwarnings.filterwarnings('ignore')\n\nclass PatchGeneratorTrainer:\n    def __init__(self, data_list, model_path=\"/kaggle/input/qwen2.5/transformers/0.5b/1\", max_length=512):\n        \"\"\"\n        Initialize the PatchGenerator with data and model configuration.\n        \n        Args:\n            data_list: List of dictionaries containing the training data\n            model_path: Path to the pre-trained model\n            max_length: Maximum sequence length for tokenization\n        \"\"\"\n        self.data = data_list\n        self.model_path = model_path\n        self.max_length = max_length\n        self.setup_logging()\n        \n    def setup_logging(self):\n        \"\"\"Setup logging configuration\"\"\"\n        logging.basicConfig(\n            level=logging.INFO,\n            format='%(asctime)s - %(levelname)s - %(message)s'\n        )\n        self.logger = logging.getLogger(__name__)\n\n    def create_prompt(self, issue_data):\n        \"\"\"Generate a structured prompt for the model\"\"\"\n        prompt = f\"\"\"Below is a Python code issue that needs to be fixed.\n\nError Type: {issue_data['problem']['error_type']}\n\nProblem Description:\n{issue_data['problem']['description']}\n\nThe patch should fix the code issue while being minimal and focused.\n\nGenerate the patch in git diff format:\"\"\"\n\n        return prompt\n    def prepare_data(self, data_list):\n        \"\"\"Prepare dataset from list of examples\"\"\"\n        training_pairs = []\n        \n        for item in data_list:\n            try:\n                prompt = self.create_prompt(item)\n                patch = item['solution']['raw_patch']\n                \n                # Only include examples with valid patches\n                if patch and patch.strip():\n                    training_pairs.append({\n                        'input_text': prompt,\n                        'output_text': patch\n                    })\n            except KeyError as e:\n                self.logger.warning(f\"Skipping malformed data entry: {e}\")\n                continue\n                \n        return training_pairs\n\n    def load_model_and_tokenizer(self):\n        \"\"\"Load and configure the model and tokenizer\"\"\"\n        self.logger.info(\"Loading tokenizer and model...\")\n        \n        try:\n            # Load model configuration\n            config = AutoConfig.from_pretrained(\n                self.model_path,\n                trust_remote_code=True,\n                local_files_only=True\n            )\n            \n            # Load tokenizer\n            self.tokenizer = AutoTokenizer.from_pretrained(\n                self.model_path,\n                local_files_only=True,\n                use_fast=True,\n                trust_remote_code=True\n            )\n            \n            # Configure quantization\n            quantization_config = BitsAndBytesConfig(\n                load_in_4bit=True,\n                bnb_4bit_compute_dtype=torch.float16,\n                bnb_4bit_use_double_quant=True,\n                bnb_4bit_quant_type='nf4'\n            )\n            \n            # Load pre-trained model\n            self.model = AutoModelForCausalLM.from_pretrained(\n                self.model_path,\n                config=config,\n                device_map=\"auto\",\n                local_files_only=True,\n                trust_remote_code=True,\n                use_safetensors=True,\n                torch_dtype=torch.float16,\n                quantization_config=quantization_config,\n                low_cpu_mem_usage=True,\n                offload_folder=\"offload\",\n                offload_state_dict=True,\n                max_memory={0: \"12GB\", \"cpu\": \"24GB\"}\n            )\n            \n            # Enable gradient checkpointing and prepare for training\n            self.model.gradient_checkpointing_enable()\n            self.model = prepare_model_for_kbit_training(self.model)\n            \n            # Apply PEFT using LoRA\n            peft_config = LoraConfig(\n                task_type=TaskType.CAUSAL_LM,\n                inference_mode=False,\n                r=8,\n                lora_alpha=32,\n                lora_dropout=0.1,\n                target_modules=[\n                    \"self_attn.q_proj\", \n                    \"self_attn.k_proj\", \n                    \"self_attn.v_proj\", \n                    \"self_attn.o_proj\",\n                    \"mlp.gate_proj\", \n                    \"mlp.up_proj\", \n                    \"mlp.down_proj\"\n                ]\n            )\n            \n            self.model = get_peft_model(self.model, peft_config)\n            self.model.print_trainable_parameters()\n            \n            self.logger.info(\"Model loaded successfully with PEFT!\")\n            \n            # Clear GPU memory if available\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n                self.logger.info(f\"GPU memory after loading: {torch.cuda.memory_allocated() / 1024**3:.2f} GB\")\n            \n        except Exception as e:\n            self.logger.error(f\"Error during model initialization: {str(e)}\")\n            raise e\n\n    def create_datasets(self, data_pairs, test_size=0.3, seed=42):\n        \"\"\"Create and split datasets for training and testing\"\"\"\n        train_pairs, test_pairs = train_test_split(\n            data_pairs, \n            test_size=test_size, \n            random_state=seed\n        )\n        \n        def preprocess_function(examples):\n            # Tokenize inputs\n            model_inputs = self.tokenizer(\n                examples['input_text'],\n                padding=True,\n                truncation=True,\n                max_length=self.max_length,\n                return_tensors=\"pt\"\n            )\n            \n            # Tokenize outputs\n            labels = self.tokenizer(\n                examples['output_text'],\n                padding=True,\n                truncation=True,\n                max_length=self.max_length,\n                return_tensors=\"pt\"\n            )\n            \n            model_inputs[\"labels\"] = labels[\"input_ids\"]\n            return model_inputs\n        \n        # Create datasets\n        train_dataset = Dataset.from_dict({\n            \"input_text\": [pair[\"input_text\"] for pair in train_pairs],\n            \"output_text\": [pair[\"output_text\"] for pair in train_pairs]\n        })\n        \n        test_dataset = Dataset.from_dict({\n            \"input_text\": [pair[\"input_text\"] for pair in test_pairs],\n            \"output_text\": [pair[\"output_text\"] for pair in test_pairs]\n        })\n        \n        # Process datasets\n        self.train_dataset = train_dataset.map(\n            preprocess_function,\n            batched=True,\n            remove_columns=train_dataset.column_names,\n            desc=\"Preprocessing training data\"\n        )\n        \n        self.test_dataset = test_dataset.map(\n            preprocess_function,\n            batched=True,\n            remove_columns=test_dataset.column_names,\n            desc=\"Preprocessing evaluation data\"\n        )\n        \n        self.test_pairs = test_pairs\n        return self.train_dataset, self.test_dataset\n\n    def train(self, output_dir=\"/kaggle/working/qwen-patch-generator\"):\n        \"\"\"Train the model\"\"\"\n        training_args = TrainingArguments(\n            output_dir=output_dir,\n            num_train_epochs=150,\n            per_device_train_batch_size=4,\n            gradient_accumulation_steps=4,\n            warmup_ratio=0.1,\n            learning_rate=1e-5,\n            logging_dir=\"./logs\",\n            fp16=True,\n            logging_steps=1,\n            save_strategy=\"epoch\",\n            save_total_limit=1,\n            evaluation_strategy=\"no\",\n            remove_unused_columns=True,\n            report_to=\"none\"\n        )\n        \n        # Define data collator\n        data_collator = DataCollatorForLanguageModeling(\n            tokenizer=self.tokenizer,\n            mlm=False\n        )\n        \n        trainer = Trainer(\n            model=self.model,\n            args=training_args,\n            train_dataset=self.train_dataset,\n            eval_dataset=self.test_dataset,\n            data_collator=data_collator,\n        )\n        \n        try:\n            self.logger.info(\"Starting training...\")\n            trainer.train()\n            \n            final_model_path = f\"{output_dir}/final\"\n            self.model.save_pretrained(final_model_path)\n            self.tokenizer.save_pretrained(final_model_path)\n            self.logger.info(f\"Final model saved to {final_model_path}\")\n            \n            # Clear memory\n            gc.collect()\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n            \n            return trainer\n            \n        except Exception as e:\n            self.logger.error(f\"Error during training: {str(e)}\")\n            raise e\n\n    def evaluate_model(self, similarity_threshold=0.7):\n        \"\"\"Evaluate model performance on test set\"\"\"\n        self.logger.info(\"Starting model evaluation...\")\n        \n        similarities = []\n        successes = 0\n        total = len(self.test_pairs)\n        \n        for idx, test_pair in enumerate(self.test_pairs):\n            try:\n                generated_patch = self.generate_patch(test_pair['input_text'])\n                similarity = self.evaluate_patch(generated_patch, test_pair['output_text'])\n                similarities.append(similarity)\n                \n                if similarity >= similarity_threshold:\n                    successes += 1\n                \n                if (idx + 1) % 10 == 0:\n                    self.logger.info(f\"Evaluated {idx + 1}/{total} test cases...\")\n                \n            except Exception as e:\n                self.logger.error(f\"Error evaluating test case {idx}: {e}\")\n                continue\n        \n        avg_similarity = np.mean(similarities)\n        accuracy = successes / total\n        \n        results = {\n            'average_similarity': avg_similarity,\n            'accuracy': accuracy,\n            'total_test_cases': total,\n            'successful_generations': successes\n        }\n        \n        self.logger.info(f\"Evaluation Results: {results}\")\n        return results\n\n    def generate_patch(self, prompt):\n        \"\"\"Generate a patch for a given problem\"\"\"\n        inputs = self.tokenizer(prompt, return_tensors=\"pt\").to(self.model.device)\n        \n        outputs = self.model.generate(\n            **inputs,\n            max_length=self.max_length,\n            num_return_sequences=1,\n            temperature=0.7,\n            top_p=0.95,\n            do_sample=True\n        )\n        \n        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n\n    def evaluate_patch(self, generated_patch, reference_patch):\n        \"\"\"Evaluate the similarity between generated and reference patches\"\"\"\n        print(\"Generated Patch\")\n        print(generated_patch)\n        print('---------------------------------------\\n')\n        similarity = SequenceMatcher(None, generated_patch, reference_patch).ratio()\n        return similarity\n\n\ndef main():\n    # Load and process the data\n    print(\"Loading and processing data...\")\n    processed_data = load_and_process_data()\n    \n    if not processed_data:\n        print(\"Error: No data processed. Exiting.\")\n        return\n    \n    # Initialize trainer\n    trainer = PatchGeneratorTrainer(processed_data)\n    \n    # Prepare data for training\n    print(\"\\nPreparing data for training...\")\n    data_pairs = trainer.prepare_data(processed_data)\n    print(f\"Total training pairs created: {len(data_pairs)}\")\n    \n    # Load model and tokenizer\n    print(\"\\nLoading model and tokenizer...\")\n    trainer.load_model_and_tokenizer()\n    \n    # Create datasets\n    print(\"\\nCreating train/test splits...\")\n    train_dataset, test_dataset = trainer.create_datasets(data_pairs)\n    print(f\"Train set size: {len(train_dataset)}\")\n    print(f\"Test set size: {len(test_dataset)}\")\n    \n    # Train model\n    print(\"\\nStarting training...\")\n    trained_trainer = trainer.train()\n    \n    # Evaluate model\n    print(\"\\nEvaluating model...\")\n    results = trainer.evaluate_model()\n    \n    print(\"\\nFinal Results:\")\n    print(f\"Average Similarity: {results['average_similarity']:.2f}\")\n    print(f\"Accuracy (>= 0.7 similarity): {results['accuracy']:.2f}\")\n    print(f\"Total Test Cases: {results['total_test_cases']}\")\n    print(f\"Successful Generations: {results['successful_generations']}\")\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-09T15:39:41.437319Z","iopub.execute_input":"2025-01-09T15:39:41.43775Z","iopub.status.idle":"2025-01-09T15:49:16.151555Z","shell.execute_reply.started":"2025-01-09T15:39:41.437708Z","shell.execute_reply":"2025-01-09T15:49:16.150574Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}