{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84795,"databundleVersionId":10821240,"sourceType":"competition"},{"sourceId":166218,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":141432,"modelId":164048}],"dockerImageVersionId":30823,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"reading_prompt = \"\"\"\nYou are analyzing a code repository to fix an issue.\n\nProblem: {problem_statement}\n\nRepository structure:\n{directory_string}\n\nPlease identify the relevant files to inspect in XML format:\n<root>\n    <entry>\n        <filepath>EXACT_PATH_FROM_REPOSITORY</filepath>\n        <strings_to_search>\n            <string_to_search>RELEVANT_CODE_PATTERN</string_to_search>\n        </strings_to_search>\n    </entry>\n</root>\n\"\"\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install -q torch transformers\n!pip install -q peft bitsandbytes\n!pip install datasets\n# # !pip install -q accelerate\n# !pip install -q vllm\n# # !pip install -q pandas pyarrow\n\nimport os\n\nimport torch\n\n# os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\n# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n# torch.cuda.set_per_process_memory_fraction(0.8, device=torch.device(\"cuda:0\"))\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import io\nimport os\nimport shutil\nimport tempfile\nimport warnings\nimport re\nfrom typing import Optional, Dict, List\nfrom dataclasses import dataclass\nimport tarfile\n\nimport pandas as pd\nfrom datasets import Dataset\nimport torch\nfrom transformers import (\n    AutoModelForCausalLM, \n    AutoTokenizer,\n    TrainingArguments,\n    Trainer\n)\nfrom peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n\nwarnings.simplefilter('ignore')\n\n# Environment settings\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\nos.environ[\"WANDB_DISABLED\"] = \"true\"\n\n# Constants\nREPO_PATH = \"repo\"\nMAX_MODEL_LEN = 32_768\n\nif os.getenv('KAGGLE_KERNEL_RUN_TYPE') or os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    model_path = '/kaggle/input/qwen2.5/transformers/0.5b/1'\nelse:\n    model_path = '/kaggle/input/qwen2.5/transformers/0.5b/1'","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# reading_prompt = \"\"\"\n# You are analyzing a code repository to fix an issue. Based on the error message and problem statement, you need to identify the relevant files that need to be modified.\n\n# Problem: {problem_statement}\n\n# Repository structure:\n# {directory_string}\n\n# Based on the error traceback and file structure above, provide the exact filepath and relevant code patterns to search in this XML format:\n# <root>\n#     <entry>\n#         <filepath>repo/astroid/nodes/node_classes.py</filepath>\n#         <strings_to_search>\n#             <string_to_search>_infer_from_values</string_to_search>\n#             <string_to_search>_infer</string_to_search>\n#         </strings_to_search>\n#     </entry>\n# </root>\n\n# Note: \n# - Use the exact filepath from the repository structure\n# - Include specific code patterns mentioned in the error trace\n# - You can add multiple entries if multiple files need to be modified\n\n# Please do not print the entire content of the file.\n# \"\"\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"reading_prompt = \"\"\"\nYou will be implementing a git diff patch to solve an issue with the code repository.\nYou will first need to select files in the file directory.\n\nThis is the problem statement.\n\n{problem_statement}\n\nThis is the file directory\n\n<directory>\n{directory_string}\n</directory>\n\nWhich files should be inspected so that we can solve the problem?\nWhen we inspect each file, what strings should be searched?\n\nReturn the strings to search in this format\n\n(explanation)\n\n<root>\n    <entry>\n        <filepath>filepath</filepath>  \n        <strings_to_search>\n            <string_to_search>_infer_from_values</string_to_search>\n            ...\n            <string_to_search>_infer_from_values</string_to_search>\n        </strings_to_search>\n    </entry>\n</root>\n\n(explanation)\n\n<root>\n    <entry>\n        <filepath>filepath</filepath>\n        <strings_to_search>\n            <string_to_search>_infer_from_values</string_to_search>\n            ...\n            <string_to_search>_infer_from_values</string_to_search>\n        </strings_to_search>\n    </entry>\n</root>\n...\n\nNotes:\n- Make sure to encode each entry between <root> and </root>\n- Return the FULL filepath - exactly as specified in <directory> and </directory>\n    - Example: repo/path/to/directory/file.py\n- If you are searching for a word instead of a substring, maybe add spaces or brackets before and after the string\n    For example, if you are searching for uses of the function `calculate`, use ` calculate(` as the search string instead\n- Prefer searching longer strings\n- Do not inspect more than 5 files\n- Only inspect the necessary files\n\"\"\".strip()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# reading_prompt = \"\"\"\n# You will be implementing a git diff patch to solve an issue with the code repository.\n# You will first need to select files in the file directory.\n\n# This is the problem statement.\n\n# {problem_statement}\n\n# This is the file directory\n\n# <directory>\n# {directory_string}\n# </directory>\n\n# Which files should be inspected so that we can solve the problem?\n# When we inspect each file, what strings should be searched?\n\n# Return the strings to search in this format:\n\n# <root>\n#     <entry>\n#         <filepath>filepath</filepath>  \n#         <strings_to_search>\n#             <string_to_search>string_to_search</string_to_search>\n#         </strings_to_search>\n#     </entry>\n# </root>\n# \"\"\".strip()\n\n# patching_prompt = \"\"\"\n# You will be implementing a git diff patch to solve an issue with the code repository.\n# This is the problem statement.\n\n# {problem_statement}\n\n# These are the files that is thought to be relevant\n\n# {file_content_string}\n\n# Write a git diff within <patch> and </patch> that fixes the problem.\n# \"\"\".strip()\n\npatching_prompt = \"\"\"\nYou will be implementing a git diff patch to solve an issue with the code repository.\nThis is the problem statement:\n\n{problem_statement}\n\nThese are the relevant files and their contents:\n\n{file_content_string}\n\nThe error occurs when a None value is being formatted with a format specifier. \nThe patch should add a check to handle None values before formatting.\n\nWrite a git diff within <patch> and </patch> that fixes the problem.\n\nExample format:\n<patch>\ndiff --git a/repo/astroid/nodes/node_classes.py b/repo/astroid/nodes/node_classes.py\n--- a/repo/astroid/nodes/node_classes.py\n+++ b/repo/astroid/nodes/node_classes.py\n@@ -4694,6 +4694,8 @@ def _infer(\n     def _infer_from_values(self, context):\n         values = self._infer(context)\n         value = next(values)\n+        if value.value is None:\n+            raise TypeError(\"Cannot apply format specifier to None value\")\n         formatted = format(value.value, format_spec.value)\n</patch>\n\"\"\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"patching_prompt = \"\"\"\nYou will be implementing a git diff patch to solve an issue with the code repository.\nThis is the problem statement.\n\n{problem_statement}\n\nThese are the files that is thought to be relevant\n\n{file_content_string}\n\nWrite a git diff within <patch> and </patch> that fixes the problem.\n\nExample:\n\n<patch>\n--- a/first.txt\n+++ b/first.txt\n@@ -1,3 +1,3 @@\n start\n-first change\n+new first change\n middle\n@@ -7,4 +7,4 @@\n some content\n-second change\n+new second change\n more content\n--- a/second.txt\n+++ b/second.txt\n@@ -1,3 +1,3 @@\n beginning\n-old line\n+new line\n end\n</patch>\n\"\"\".strip()\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def stringify_directory(directory):\n    full_paths = []\n    for root, dirs, files in os.walk(directory):\n        for file in files:\n            full_path = os.path.join(root, file)\n            full_paths.append(full_path)\n    return \"\\n\".join(full_paths)\n\ndef extract_file_query(xml_content):\n\n    if not xml_content or \"<root>\" not in xml_content:\n        print(\"Invalid XML response from model\")\n        return {}\n    parsed_data = {}\n    pattern = r'<root>(.*?)</root>'\n    matches = re.findall(pattern, xml_content, re.DOTALL)\n    \n    for match in matches:\n        try:\n            import xml.etree.ElementTree as ET\n            root = ET.fromstring(\"<root>\" + match + \"</root>\")\n            \n            for entry in root.findall(\"entry\"):\n                filepath = entry.find(\"filepath\")\n                filepath_text = filepath.text.strip() if filepath is not None else None\n                \n                strings_container = entry.find(\"strings_to_search\")\n                search_strings = []\n                if strings_container is not None:\n                    for s in strings_container.findall(\"string_to_search\"):\n                        if s.text is not None:\n                            search_strings.append(s.text.strip())\n                \n                parsed_data[filepath_text] = search_strings\n        except:\n            print(\"Error parsing output\", xml_content)\n            return \"\"\n        \n    return parsed_data\n\ndef extract_patch_string(text):\n    pattern = r'<patch>(.*?)</patch>'\n    matches = re.findall(pattern, text, re.DOTALL)\n    if not matches:\n        return None\n    return \"\\n\".join(matches)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def fetch_file_contents(files_to_search, context_lines=10):\n\n    if not files_to_search or not isinstance(files_to_search, dict):\n        print(\"Invalid file query structure\")\n        return \"\"\n    def find_lines_in_files_with_context(search_map, context_lines):\n        all_matches_per_file = []\n\n        for path, terms in search_map.items():\n            if not os.path.isfile(path):\n                all_matches_per_file.append([])\n                continue\n\n            with open(path, \"r\", encoding=\"utf-8\", errors=\"replace\") as f:\n                lines = f.readlines()\n\n            file_snippets = []\n            num_lines = len(lines)\n\n            for i, line in enumerate(lines, start=1):\n                if any(t in line for t in terms):\n                    start_idx = max(1, i - context_lines)\n                    end_idx = min(num_lines, i + context_lines)\n                    snippet = []\n                    for snippet_no in range(start_idx, end_idx + 1):\n                        text_content = lines[snippet_no - 1].rstrip(\"\\n\")\n                        snippet.append((snippet_no, text_content))\n                    file_snippets.append(snippet)\n\n            all_matches_per_file.append(file_snippets)\n\n        return all_matches_per_file\n\n    snippets = find_lines_in_files_with_context(files_to_search, context_lines)\n    output = []\n    \n    for filepath, file_snippets in zip(files_to_search.keys(), snippets):\n        output.append(f\"FILE: {filepath[len(REPO_PATH) + 1:]}\")\n        if not file_snippets:\n            output.append(\"No matches found.\")\n        else:\n            for snippet in file_snippets:\n                start_line = snippet[0][0]\n                end_line = snippet[-1][0]\n                output.append(f\"Lines {start_line} to {end_line}:\")\n                for line_no, text in snippet:\n                    output.append(f\"{line_no:3d} | {text}\")\n                output.append(\"\")\n        output.append(\"=\" * 60)\n        output.append(\"\")\n    \n    return \"\\n\".join(output)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install --upgrade transformers\n\nfrom transformers import AutoModelForCausalLM, BitsAndBytesConfig\n\n\nclass GitPatchModel:\n    def __init__(self, model_path: str):  # Fixed double underscores\n        self.tokenizer = AutoTokenizer.from_pretrained(model_path)\n\n        ## Adding new\n        config = AutoConfig.from_pretrained(\n                self.model_path,\n                trust_remote_code=True,\n                local_files_only=True\n            )\n        \n        # 4-bit quantization setup\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            llm_int8_enable_fp32_cpu_offload=True\n\n            \n        )\n        \n        # Load base model\n        print(\"Loading base model...\")\n\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        \n        self.model = AutoModelForCausalLM.from_pretrained(\n            model_path,\n            device_map=\"auto\",\n            trust_remote_code=True,\n            quantization_config=quantization_config,\n            torch_dtype=torch.float16,\n            use_cache = False,\n            low_cpu_mem_usage=True,\n            offload_folder=\"offload\",\n            offload_state_dict=True,\n            max_memory={0: \"8GB\", \"cpu\": \"16GB\"}\n        )\n        \n        # Prepare for training\n        self.model = prepare_model_for_kbit_training(self.model)\n        \n        # LoRA config\n        lora_config = LoraConfig(\n            r=8,\n            lora_alpha=32,\n            target_modules=[\"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            lora_dropout=0.1,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\"\n        )\n        \n        self.model = get_peft_model(self.model, lora_config)\n        \n    def generate(self, prompt: str, max_length: int = 512) -> str:\n        inputs = self.tokenizer(prompt, return_tensors=\"pt\").to(self.model.device)\n        outputs = self.model.generate(\n            **inputs,\n            max_tokens = 512,\n            #max_length=max_length,\n            temperature=1.0,\n            do_sample=True,\n            pad_token_id=self.tokenizer.pad_token_id\n        )\n        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n\n    def collate_fn(self, examples):\n        \"\"\"Collate function for batching examples.\"\"\"\n        # Get prompts and completions\n        prompts = [ex[\"prompt\"] for ex in examples]\n        completions = [ex[\"completion\"] for ex in examples]\n        \n        # First tokenize both inputs with the same max length\n        input_encodings = self.tokenizer(\n            prompts,\n            padding=True,\n            truncation=True,\n            max_length=MAX_MODEL_LEN,\n            return_tensors=\"pt\"\n        ).to('cpu')\n        \n        label_encodings = self.tokenizer(\n            completions,\n            padding=True,\n            truncation=True,\n            max_length=MAX_MODEL_LEN,\n            return_tensors=\"pt\"\n        ).to('cpu')\n    \n        # Create labels tensor with -100 for padding tokens\n        labels = label_encodings[\"input_ids\"].clone()\n        labels[label_encodings[\"attention_mask\"] == 0] = -100\n        \n        # Ensure input_ids and labels have same size\n        input_length = input_encodings[\"input_ids\"].size(1)\n        label_length = labels.size(1)\n\n        # print(\"input_ids shape:\", input_encodings[\"input_ids\"].shape)\n        # print(\"labels shape:\", labels.shape)\n        \n        # If sizes don't match, pad the shorter one\n        if input_length > label_length:\n            # Pad labels\n            padding_size = input_length - label_length\n            labels = torch.nn.functional.pad(labels, (0, padding_size), value=-100)\n        elif label_length > input_length:\n            # Pad inputs\n            padding_size = label_length - input_length\n            input_encodings[\"input_ids\"] = torch.nn.functional.pad(input_encodings[\"input_ids\"], (0, padding_size), value=self.tokenizer.pad_token_id)\n            input_encodings[\"attention_mask\"] = torch.nn.functional.pad(input_encodings[\"attention_mask\"], (0, padding_size), value=0)\n    \n        # Create final batch\n        batch = {\n            \"input_ids\": input_encodings[\"input_ids\"],\n            \"attention_mask\": input_encodings[\"attention_mask\"],\n            \"labels\": labels\n        }\n    \n        # Verify shapes match\n        assert batch[\"input_ids\"].size() == batch[\"labels\"].size(), \\\n            f\"Input shape {batch['input_ids'].size()} != Labels shape {batch['labels'].size()}\"\n        \n        return batch\n\n\nclass GitPatchModel:\n    def __init__(self, model_path: str):\n        \"\"\"Initialize the model with memory optimizations.\"\"\"\n        self.model_path = model_path  # Store model_path as instance variable\n        self.tokenizer = AutoTokenizer.from_pretrained(model_path)\n        \n        # Configure tokenizer\n        self.tokenizer.pad_token = self.tokenizer.eos_token\n        self.tokenizer.padding_side = 'right'\n\n        # Load configuration\n        config = AutoConfig.from_pretrained(\n            self.model_path,\n            trust_remote_code=True,\n            local_files_only=True\n        )\n        \n        # 4-bit quantization setup for memory efficiency\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            llm_int8_enable_fp32_cpu_offload=True\n        )\n        \n        print(\"Loading base model with memory optimizations...\")\n        self.model = AutoModelForCausalLM.from_pretrained(\n            model_path,\n            device_map=\"auto\",\n            trust_remote_code=True,\n            quantization_config=quantization_config,\n            torch_dtype=torch.float16,\n            use_cache=False,\n            low_cpu_mem_usage=True,\n            offload_folder=\"offload\",\n            offload_state_dict=True,\n            max_memory={0: \"4GB\", \"cpu\": \"8GB\"}\n        )\n        \n        # Prepare for training with LoRA\n        print(\"Preparing model for training...\")\n        self.model = prepare_model_for_kbit_training(self.model)\n        \n        # LoRA config for memory-efficient fine-tuning\n        lora_config = LoraConfig(\n            r=8,\n            lora_alpha=32,\n            target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"],\n            lora_dropout=0.1,\n            bias=\"none\",\n            task_type=\"CAUSAL_LM\"\n        )\n\n        \n        \n        self.model = get_peft_model(self.model, lora_config)\n        \n    # def collate_fn(self, examples):\n    #     \"\"\"Collate function for batching examples.\"\"\"\n    #     prompts = [ex[\"prompt\"] for ex in examples]\n    #     completions = [ex[\"completion\"] for ex in examples]\n        \n    #     max_length = 512\n        \n    #     # Ensure tensors stay on CPU\n    #     inputs = self.tokenizer(\n    #         prompts,\n    #         padding=True,\n    #         truncation=True,\n    #         max_length=max_length,\n    #         return_tensors=\"pt\"\n    #     ).to('cpu')  # Explicitly move to CPU\n        \n    #     labels = self.tokenizer(\n    #         completions,\n    #         padding=True,\n    #         truncation=True,\n    #         max_length=max_length,\n    #         return_tensors=\"pt\"\n    #     ).to('cpu')  # Explicitly move to CPU\n        \n    #     batch = {\n    #         \"input_ids\": inputs[\"input_ids\"],\n    #         \"attention_mask\": inputs[\"attention_mask\"],\n    #         \"labels\": labels[\"input_ids\"].clone()\n    #     }\n    #     print(len(labels))\n    #     print(len(inputs))\n    #     # Replace padding tokens\n    #     batch[\"labels\"][labels[\"attention_mask\"] == 0] = -100\n        \n    #     # Ensure all tensors are on CPU\n    #     batch = {k: v.to('cpu') for k, v in batch.items()}\n        \n    #     return batch\n\n    def generate(self, prompt: str, max_length: int = 512) -> str:\n        inputs = self.tokenizer(prompt, return_tensors=\"pt\").to(self.model.device)\n        outputs = self.model.generate(\n            **inputs,\n            #max_length=max_length,\n            max_new_tokens=512,\n            temperature=1.0,\n            do_sample=True,\n            pad_token_id=self.tokenizer.pad_token_id\n        )\n        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n\n    def collate_fn(self, examples):\n        \"\"\"\n        Collate function for batching examples.\n        Expects examples to be a list of dicts with keys \"prompt\" and \"completion\".\n        \"\"\"\n        # Gather prompts and completions\n        prompts = [ex[\"prompt\"] for ex in examples]\n        completions = [ex[\"completion\"] for ex in examples]\n        \n        # Set a fixed max_length for both\n        max_length = 512\n        \n        # 1) Tokenize the prompts\n        input_encodings = self.tokenizer(\n            prompts,\n            padding=\"max_length\",   # Ensures uniform length\n            truncation=True,\n            max_length=max_length,\n            return_tensors=\"pt\"\n        )\n        \n        # 2) Tokenize the completions\n        label_encodings = self.tokenizer(\n            completions,\n            padding=\"max_length\",   # Ensures uniform length\n            truncation=True,\n            max_length=max_length,\n            return_tensors=\"pt\"\n        )\n        \n        # 3) Create labels by cloning the label_encodings input_ids\n        labels = label_encodings[\"input_ids\"].clone()\n        \n        # Replace all padding token indices with -100 so they are ignored by the loss\n        labels[label_encodings[\"attention_mask\"] == 0] = -100\n        \n        # (Optional) Move everything to CPU here (depending on your setup):\n        input_ids = input_encodings[\"input_ids\"].to(\"cpu\")\n        attention_mask = input_encodings[\"attention_mask\"].to(\"cpu\")\n        labels = labels.to(\"cpu\")\n        \n        # Debug prints: confirm shapes match\n        # print(f\"input_ids.shape: {input_ids.shape}\")\n        # print(f\"attention_mask.shape: {attention_mask.shape}\")\n        # print(f\"labels.shape: {labels.shape}\")\n        \n        # 4) Wrap in a dict for Trainer\n        batch = {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n            \"labels\": labels\n        }\n        \n        return batch\n\n\ndef train_model(model_path: str, data_path: str, repos_dir: str, output_dir: str):\n    \"\"\"Train the model on the prepared dataset.\"\"\"\n    print(\"Initializing model and data processor...\")\n    model = GitPatchModel(model_path)\n    data_processor = GitPatchDataProcessor(data_path, repos_dir)\n    \n    print(\"Preparing training data...\")\n    train_examples = data_processor.prepare_training_examples()\n    if not train_examples:\n        raise ValueError(\"No training examples generated!\")\n    \n    train_dataset = Dataset.from_list(train_examples)\n    print(f\"Created dataset with {len(train_dataset)} examples\")\n    \n    training_args = TrainingArguments(\n            output_dir=output_dir,\n            num_train_epochs=500,\n            per_device_train_batch_size=1,  # Reduced batch size\n            gradient_accumulation_steps=8,\n            learning_rate=2e-4,\n            fp16=True,\n            save_strategy=\"epoch\",\n            save_total_limit=1,\n            logging_steps=10,\n            remove_unused_columns=False,\n            report_to=\"none\",\n            gradient_checkpointing=True,\n            max_grad_norm=0.3,\n            # max_length=2048,  # Add explicit max length\n            # pad_to_max_length=True,  # Ensure padding\n            dataloader_drop_last=True  # Drop last incomplete batch\n        )\n    \n    print(\"Initializing trainer...\")\n    trainer = Trainer(\n        model=model.model,\n        args=training_args,\n        train_dataset=train_dataset,\n        data_collator=model.collate_fn\n    )\n    \n    print(\"Starting training...\")\n    trainer.train()\n    \n    print(f\"Saving model to {output_dir}...\")\n    trainer.save_model(output_dir)\n    print(\"Training completed successfully!\")\n\ndef predict(model: GitPatchModel, problem_statement: str, repo_archive) -> str:\n    \"\"\"Generate a patch prediction for a given problem.\"\"\"\n    # Extract repository to temporary directory\n    with tempfile.TemporaryDirectory() as tmpdir:\n        repo_path = os.path.join(tmpdir, 'repo')\n        os.makedirs(repo_path)\n        \n        # Extract repo archive\n        repo_archive.seek(0)\n        with tarfile.open(fileobj=repo_archive, mode='r:*') as tar:\n            tar.extractall(repo_path)\n        \n        # Get directory structure\n        directory_string = stringify_directory(repo_path)\n        \n        # Create prompt\n        prompt = (\n            f\"Problem: {problem_statement}\\n\\n\"\n            f\"Repository structure:\\n{directory_string}\\n\\n\"\n            \"Generate a patch to fix this issue.\"\n        )\n        \n        # Generate prediction\n        inputs = model.tokenizer(prompt, return_tensors=\"pt\")\n        outputs = model.model.generate(\n            **inputs,\n            max_length=2048,\n            num_return_sequences=1,\n            pad_token_id=model.tokenizer.eos_token_id\n        )\n        \n        prediction = model.tokenizer.decode(outputs[0], skip_special_tokens=True)\n        \n        # Extract patch from prediction\n        if \"<patch>\" in prediction and \"</patch>\" in prediction:\n            patch = prediction.split(\"<patch>\")[1].split(\"</patch>\")[0]\n        else:\n            patch = prediction\n            \n        return patch\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nfrom typing import List, Optional\nimport os\nimport glob\nfrom pathlib import Path\n\nclass GitPatchDataProcessor:\n    def __init__(self, data_path: str, repos_dir: str):\n        \"\"\"Initialize the GitPatchDataProcessor.\n        \n        Args:\n            data_path: Path to the parquet file containing the dataset\n            repos_dir: Directory containing the git repositories\n        \"\"\"\n        # Load and validate the parquet file\n        if not os.path.exists(data_path):\n            raise FileNotFoundError(f\"Data file not found: {data_path}\")\n        self.data_df = pd.read_parquet(data_path)\n        \n        # Validate repos directory\n        if not os.path.exists(repos_dir):\n            raise FileNotFoundError(f\"Repos directory not found: {repos_dir}\")\n        self.repos_dir = repos_dir\n        \n        # Validate required columns\n        required_columns = ['instance_id', 'problem_statement', 'patch']\n        missing_columns = [col for col in required_columns if col not in self.data_df.columns]\n        if missing_columns:\n            raise ValueError(f\"Missing required columns in dataset: {missing_columns}\")\n            \n    def stringify_directory(self, repo_path: str) -> str:\n        \"\"\"Convert directory structure to a string representation.\n        \n        Args:\n            repo_path: Path to the repository\n            \n        Returns:\n            String representation of the directory structure\n        \"\"\"\n        if not os.path.exists(repo_path):\n            print(f\"Warning: Repository path does not exist: {repo_path}\")\n            return \"\"\n            \n        result = []\n        for root, dirs, files in os.walk(repo_path):\n            # Skip .git directory\n            if '.git' in dirs:\n                dirs.remove('.git')\n                \n            level = root.replace(repo_path, '').count(os.sep)\n            indent = '  ' * level\n            result.append(f'{indent}{os.path.basename(root)}/')\n            \n            for file in files:\n                result.append(f'{indent}  {file}')\n                \n        return '\\n'.join(result)\n        \n    def read_file_content(self, file_path: str) -> str:\n        \"\"\"Read content of a file safely.\n        \n        Args:\n            file_path: Path to the file to read\n            \n        Returns:\n            Content of the file or empty string if file cannot be read\n        \"\"\"\n        try:\n            with open(file_path, 'r', encoding='utf-8') as f:\n                return f.read()\n        except Exception as e:\n            print(f\"Warning: Could not read file {file_path}: {str(e)}\")\n            return \"\"\n            \n    def prepare_training_examples(self, \n                                start_idx: int = 0, \n                                end_idx: Optional[int] = None) -> List[dict]:\n        \"\"\"Prepare training examples from the dataset.\n        \n        Args:\n            start_idx: Starting index in the dataset\n            end_idx: Ending index in the dataset (exclusive)\n            \n        Returns:\n            List of training examples\n        \"\"\"\n        examples = []\n        end_idx = end_idx or len(self.data_df)\n        \n        print(f\"Preparing examples from index {start_idx} to {end_idx}\")\n        print(f\"Total rows in dataset: {len(self.data_df)}\")\n        \n        for idx in range(start_idx, end_idx):\n            try:\n                # Get basic information\n                problem_statement = self.data_df[\"problem_statement\"][idx]\n                instance_id = self.data_df['instance_id'][idx]\n                repo_path = os.path.join(self.repos_dir, f\"repo__{instance_id}\")\n                \n                print(f\"\\nProcessing example {idx}:\")\n                print(f\"Instance ID: {instance_id}\")\n                print(f\"Repo path: {repo_path}\")\n                \n                # Get directory structure\n                directory_string = self.stringify_directory(repo_path)\n                if not directory_string:\n                    print(f\"Warning: Empty directory structure for {repo_path}\")\n                    continue\n                    \n                # Create file navigation example\n                nav_prompt = (\n                    f\"Problem: {problem_statement}\\n\\n\"\n                    f\"Repository structure:\\n{directory_string}\\n\\n\"\n                    \"Which files should be modified to fix this issue?\"\n                )\n                \n                # Create patch generation example\n                patch_content = self.data_df['patch'][idx]\n                if not patch_content:\n                    print(f\"Warning: Empty patch content for index {idx}\")\n                    continue\n                    \n                patch_prompt = (\n                    f\"Problem: {problem_statement}\\n\\n\"\n                    f\"Repository structure:\\n{directory_string}\\n\\n\"\n                    \"Generate a patch to fix this issue.\"\n                )\n                \n                # Add both examples\n                examples.append({\n                    \"prompt\": nav_prompt,\n                    \"completion\": self.extract_files_from_patch(patch_content)\n                })\n                \n                examples.append({\n                    \"prompt\": patch_prompt,\n                    \"completion\": f\"<patch>{patch_content}</patch>\"\n                })\n                \n                print(f\"Added navigation and patch examples for index {idx}\")\n                \n            except Exception as e:\n                print(f\"Error processing index {idx}: {str(e)}\")\n                continue\n                \n        print(f\"\\nTotal examples generated: {len(examples)}\")\n        return examples\n        \n    def extract_files_from_patch(self, patch_content: str) -> str:\n        \"\"\"Extract modified file paths from a patch.\n        \n        Args:\n            patch_content: Content of the patch\n            \n        Returns:\n            List of modified files as a string\n        \"\"\"\n        files = []\n        for line in patch_content.split('\\n'):\n            if line.startswith('diff --git'):\n                # Extract the second file path (b/path/to/file)\n                parts = line.split()\n                if len(parts) >= 4:\n                    file_path = parts[3][2:]  # Remove b/ prefix\n                    files.append(file_path)\n                    \n        return '\\n'.join(files)\n        \n    def validate_repository(self, repo_path: str) -> bool:\n        \"\"\"Validate that a repository directory exists and contains files.\n        \n        Args:\n            repo_path: Path to the repository\n            \n        Returns:\n            True if repository is valid, False otherwise\n        \"\"\"\n        if not os.path.exists(repo_path):\n            print(f\"Repository directory does not exist: {repo_path}\")\n            return False\n            \n        if not os.path.isdir(repo_path):\n            print(f\"Repository path is not a directory: {repo_path}\")\n            return False\n            \n        files = glob.glob(os.path.join(repo_path, '**/*'), recursive=True)\n        if not files:\n            print(f\"Repository directory is empty: {repo_path}\")\n            return False\n            \n        return True","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def predict_inner(model: GitPatchModel, problem_statement: str, directory: str) -> str:\n#     # Get file query\n#     directory_string = stringify_directory(directory)\n#     query_prompt = reading_prompt.format(\n#         problem_statement=problem_statement,\n#         directory_string=directory_string\n#     )\n#     query_response = model.generate(query_prompt)\n#     file_query = extract_file_query(query_response)\n#     print(\"File query:\", file_query)\n\n#     if not file_query:\n#         return None\n\n#     # Get file contents\n#     file_content_string = fetch_file_contents(file_query)\n#     print(\"File contents:\", file_content_string)\n\n#     # Generate patch\n#     patch_prompt = patching_prompt.format(\n#         problem_statement=problem_statement,\n#         file_content_string=file_content_string\n#     )\n#     patch_response = model.generate(patch_prompt)\n#     patch_string = extract_patch_string(patch_response)\n#     print(\"Generated patch:\", patch_string)\n\n#     return patch_string\n\n# def predict(model: GitPatchModel, \n#             problem_statement: str, \n#             repo_archive: io.BytesIO) -> Optional[str]:\n#     # Extract repository\n#     with open('repo_archive.tar', 'wb') as f:\n#         f.write(repo_archive.read())\n        \n#     if os.path.exists(REPO_PATH):\n#         shutil.rmtree(REPO_PATH)\n#     shutil.unpack_archive('repo_archive.tar', extract_dir=REPO_PATH)\n#     os.remove('repo_archive.tar')\n\n#     try:\n#         patch_string = predict_inner(model, problem_statement, REPO_PATH)\n#     finally:\n#         shutil.rmtree(REPO_PATH)\n\n#     return patch_string","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport gc\n\ndef clear_cuda_memory():\n    \"\"\"\n    Clears CUDA memory by emptying cache and collecting garbage.\n    Should be called between large operations.\n    \"\"\"\n    # Clear pytorch's cuda cache\n    torch.cuda.empty_cache()\n    # Run garbage collector\n    gc.collect()\n\nclass MemoryManagement:\n    \"\"\"\n    Context manager for handling CUDA memory in PyTorch operations\n    \"\"\"\n    def __init__(self, model=None):\n        self.model = model\n\n    def __enter__(self):\n        clear_cuda_memory()\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        clear_cuda_memory()\n\ndef optimize_model_memory(model):\n    \"\"\"\n    Optimizes model memory usage by moving to CPU when not in use\n    \"\"\"\n    if hasattr(model, 'cpu'):\n        model.cpu()\n    clear_cuda_memory()\n    return model\n\ndef chunked_generate(model, tokenizer, prompt, chunk_size=512, max_length=512, device='cuda'):\n    \"\"\"\n    Processes generation in chunks to manage memory better\n    \"\"\"\n    tokens = tokenizer(prompt, return_tensors=\"pt\").to(device)\n    \n    # Process in chunks\n    output_tokens = []\n    for i in range(0, max_length, chunk_size):\n        with MemoryManagement():\n            current_length = min(chunk_size, max_length - i)\n            chunk_output = model.generate(\n                **tokens,\n                max_length=current_length,\n                pad_token_id=tokenizer.eos_token_id\n            )\n            output_tokens.extend(chunk_output[0][tokens.input_ids.shape[1]:])\n            \n            # Update input for next iteration\n            tokens = tokenizer(\n                tokenizer.decode(output_tokens),\n                return_tensors=\"pt\"\n            ).to(device)\n    \n    return tokenizer.decode(output_tokens)\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_inner(model: GitPatchModel, problem_statement: str, directory: str) -> str:\n    # Get file query\n    directory_string = stringify_directory(directory)\n    \n    # Use the original reading_prompt but with better instructions\n    query_prompt = reading_prompt.format(\n        problem_statement=problem_statement,\n        directory_string=directory_string\n    )\n    \n    print(\"\\nSending query prompt...\")\n    query_response = model.generate(query_prompt)\n    print(f\"\\nQuery response: {query_response}\")\n    \n    file_query = extract_file_query(query_response)\n    print(\"\\nFile query:\", file_query)\n    \n    if not file_query:\n        print(\"No valid file query extracted\")\n        return None\n\n    # Get file contents\n    file_content_string = fetch_file_contents(file_query)\n    print(\"\\nFile contents:\", file_content_string)\n\n    # Generate patch with the original patching prompt\n    patch_prompt = patching_prompt.format(\n        problem_statement=problem_statement,\n        file_content_string=file_content_string\n    )\n    \n    print(\"\\nSending patch prompt...\")\n\n    patch_response = model.generate(patch_prompt)\n    print(f\"\\nPatch response: {patch_response}\")\n    \n    patch_string = extract_patch_string(patch_response)\n    print(\"\\nGenerated patch:\", patch_string)\n    return patch_string\n\ndef predict(model: GitPatchModel, \n            problem_statement: str, \n            repo_archive: io.BytesIO) -> Optional[str]:\n    # Extract repository\n    with open('repo_archive.tar', 'wb') as f:\n        f.write(repo_archive.read())\n        \n    if os.path.exists(REPO_PATH):\n        shutil.rmtree(REPO_PATH)\n    shutil.unpack_archive('repo_archive.tar', extract_dir=REPO_PATH)\n    os.remove('repo_archive.tar')\n\n    try:\n        clear_cuda_memory()\n        optimize_model_memory(model)\n        patch_string = predict_inner(model, problem_statement, REPO_PATH)\n    finally:\n        shutil.rmtree(REPO_PATH)\n    \n    return patch_string","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!mkdir -p /kaggle/tmp/konwinski-prize-alt\n!unzip -q -o /kaggle/input/konwinski-prize/data.a_zip -d /kaggle/tmp/konwinski-prize-alt/ 2>/dev/null || true","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# torch.cuda.empty_cache()\n# torch.cuda.memory_summary()\n\n\n\nfrom transformers import (\n    AutoModelForCausalLM, \n    AutoTokenizer, \n    BitsAndBytesConfig,\n    AutoConfig\n)\nfrom peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training\n\nif __name__ == \"__main__\":\n    # Training\n    if not os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n        train_model(\n            model_path=model_path,\n            data_path='/kaggle/tmp/konwinski-prize-alt/data/data.parquet',\n            repos_dir='/kaggle/tmp/konwinski-prize-alt/data/repos',\n            output_dir='./git-patch-model'\n        )\n    \n    # Inference\n    def get_problem(problem_index):\n        df = pd.read_parquet('/kaggle/tmp/konwinski-prize-alt/data/data.parquet')\n        problem_statement = df[\"problem_statement\"][problem_index]\n        repo_path = f\"/kaggle/tmp/konwinski-prize-alt/data/repos/repo__{df['instance_id'][problem_index]}\"\n        \n        with tempfile.TemporaryDirectory() as tmpdir:\n            shutil.make_archive(os.path.join(tmpdir, 'a_repo'), 'tar', repo_path)\n            with open(os.path.join(tmpdir, 'a_repo.tar'), 'rb') as f:\n                repo_archive = io.BytesIO(f.read())\n        \n        return problem_statement, repo_path, repo_archive\n    \n    if not os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n        # Load trained model\n        model = GitPatchModel(model_path)\n        model.model.load_adapter(\"./git-patch-model\", adapter_name=\"default\")\n\n        \n        # Test on demo problemd\n        demo_problem_index = 0\n        problem_statement, repo_path, repo_archive = get_problem(demo_problem_index)\n        clear_cuda_memory()\n        optimize_model_memory(model)\n        patch_string = predict(model, problem_statement, repo_archive)\n        print(\"\\nFinal patch:\")\n        print(patch_string)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}