Set up data generation pipeline

- Add slurm script for data generation
- Finalize Qwen-2.5-coder-32B model for data generation

.gitignore

+data_generation/logs
\ No newline at end of file

data_generation/main.py

+import re
+from datasets import load_dataset, IterableDataset, concatenate_datasets
+import torch
+from transformers import AutoTokenizer
+import transformers
+import datetime
+
+
+def data_generator(dataset):
+    for row in dataset:
+        yield row
+
+start_model_loading = datetime.datetime.now()
+model = "meta-llama/Llama-3.1-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model)
+pipeline = transformers.pipeline("text-generation",
+model=model, model_kwargs={"torch_dtype": torch.bfloat16}, device_map = "auto")
+end_model_loading = datetime.datetime.now()
+print(f'Model: {model} loaded successfully in {(end_model_loading-start_model_loading).total_seconds()} seconds.')
+
+original_dataset = load_dataset("codeparrot/github-code", streaming=True, split="train", licenses=["mit", "isc"], languages = ['C#'], filter_languages=True, trust_remote_code=True)
+
+BATCH_SIZE = 100
+instance_number = 0
+batch_dataset = []
+is_dataset_created = False
+ss=0
+for example in original_dataset:
+    prompt = f"""
+    <|begin_of_text|><|start_header_id|>system<|end_header_id|>
+    You are a C# expert.
+    Your task is to refine the C# code enclosed within tags [C#] and [/C#]. 
+    Refined code should be enclosed with tags [refined_C#] and [/refined_C#].
+    Summary of changes should be enclosed with [code_changes] and [/code_changes].
+    You do not do anything more than user asks you do it.
+    You do not generate any additional text.
+    <|eot_id|><|start_header_id|>user<|end_header_id|>
+
+    You have the refine the code based on principles:
+    1. Class name: Refine C# class name by making sure:
+        a. It is PascalCase.
+        b. It is logical. By logical, I mean, it should be a noun and it should denote what it does.
+    2. Property name: Refine the properties in the class by making sure:
+        a. Private properties have _ before their names and public properties are PascalCase.
+        b. It is logical. By logical, I mean, there should be no unnecessary repetition. For example property name as 'empId' in class Employee is not clean as 'emp' is unnecessary. This is because the class name already sets the context. The correct name would be 'id'. But make sure that the property name makes sense.
+    3. Object name: Refine the objects instantiated by making sure:
+        a. It is logical. By logical I mean, the names should be expressive. For example 'Employee e' is not as expressive 'Employee employee'.
+    4. Method name: Refine methods in the class by making sure:
+        a. It is logical. By logical, I mean the name of method should express what it does. 
+        b. It is PascalCase.
+    [C#]
+    {example["code"]}
+    [/C#]<|eot_id|><|start_header_id|>assistant<|end_header_id|>[refined_C#]
+    """
+
+    no_of_retries = 0
+    should_try = True
+    while should_try:
+        sequences = pipeline(prompt,temperature=0.2,
+            top_p=0.9,
+            num_return_sequences=1,
+            eos_token_id=tokenizer.eos_token_id,
+            max_length=5000,
+            repetition_penalty=1.1)
+        
+        output = ""
+        for seq in sequences:
+            output += seq["generated_text"]
+
+        code_pattern = r'\[refined_C#\](.*?)\[/refined_C#\]'
+        code_matches = re.findall(code_pattern, output, re.DOTALL)
+        is_code_okay = False
+        if len(code_matches) == 2:
+            refined_code = code_matches[1]
+            is_code_okay = True
+
+
+        is_summary_okay = False
+        summary_pattern = r'\[code_changes\](.*?)\[/code_changes\]'
+        summary_matches = re.findall(summary_pattern, output, re.DOTALL)
+        if len(summary_matches) == 2:
+            summary = summary_matches[1]
+            is_summary_okay = True
+
+        if (is_code_okay and is_summary_okay):
+            batch_dataset.append({'code': example["code"], 'refined code': refined_code, 'summary': summary})
+            instance_number += 1
+
+        no_of_retries += 1
+        if (is_code_okay and is_summary_okay) or no_of_retries == 5:
+            should_try = False
+    
+    if instance_number == BATCH_SIZE:
+        if not is_dataset_created:
+            new_dataset = IterableDataset.from_generator(data_generator, gen_kwargs={"dataset": batch_dataset})
+            new_dataset.push_to_hub('llama-3-1-refined-code')
+            is_dataset_created = True
+            print('Pushed data for first time')
+        else:
+            refined_code_dataset = load_dataset('llama-3-1-refined-code', streaming = True)
+            new_dataset = concatenate_datasets(refined_code_dataset, IterableDataset.from_generator(data_generator, gen_kwargs={"dataset": batch_dataset}))
+            new_dataset.push_to_hub('llama-3-1-refined-code')
+            print('Pushed data again')
+
+        instance_number = 0
+        batch_dataset = []
+        ss += 1
+
+    if ss == 3:
+        break
\ No newline at end of file

data_generation/main_qwen.py

+import re
+from datasets import load_dataset, Dataset, concatenate_datasets
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import datetime
+
+
+def data_generator(dataset):
+    for row in dataset:
+        yield row
+
+SKIP_INDEX = 2000
+print(f'Starting skipping of dataset at {datetime.datetime.now()}', flush=True)
+original_dataset = load_dataset("codeparrot/github-code", streaming=True, split="train", licenses=["mit", "isc"], languages = ['C#'], filter_languages=True, trust_remote_code=True)
+original_dataset = original_dataset.skip(SKIP_INDEX)
+print(f'Skipped the dataset for {SKIP_INDEX} samples at {datetime.datetime.now()}', flush=True)
+
+start_model_loading = datetime.datetime.now()
+model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+end_model_loading = datetime.datetime.now()
+print(f'Model: {model_name} loaded successfully in {(end_model_loading-start_model_loading).total_seconds()} seconds. Current Time: {datetime.datetime.now()}', flush=True)
+
+BATCH_SIZE = 100
+TOTAL_SAMPLES = 4500
+NO_OF_BATCHES = TOTAL_SAMPLES // BATCH_SIZE
+instance_number = 0
+batch_number = 0
+batch_dataset = []
+is_dataset_created = False
+DATASET_NAME = f'qwen-refined-code-{SKIP_INDEX}'
+DATASET_NAME_WITH_USERNAME = f'atharva2721/{DATASET_NAME}'
+
+MAX_MODEL_TOKENS = 32768
+
+system_prompt = f"""
+    You are Qwen, created by Alibaba Cloud. You are a C# expert.
+    Your task is to refine the C# code enclosed within tags [C#] and [/C#]. 
+    Refined code should be enclosed with tags [refined_C#] and [/refined_C#]. It should only contain executable code and no additional text.
+    Summary of changes should be enclosed with [code_changes] and [/code_changes].
+    You do not do anything more than user asks you do it.
+    You do not generate any additional text.
+    """
+
+for example in original_dataset:
+    user_prompt = f"""
+    Refine C# Code Based on Clean Code and Design Principles.
+
+    The goal of this task is to improve the quality, readability, and maintainability of the provided C# code. Apply the following principles step by step, ensuring that the resulting code is clean, modular, and adheres to object-oriented design best practices.
+
+    You have been provided with principles and some examples wrapped inside [example_code] and [/example_code]to understand their meanings. Understand them to refine the code.
+    
+    1. Class Naming
+    a. Use PascalCase for all class names (e.g., Invoice, Employee).
+    b. Ensure class names are logical and represent a clear purpose:
+        - A class name should be a noun that describes its role or entity.
+        - Avoid vague or overly generic names like ManagerClass. Instead, use meaningful names like InvoiceManager or PayrollCalculator.
+        - Ensure the name reflects what the class does or represents.
+
+    Example:
+    Before:
+    [example_code]
+    public class ManagerClass
+    {{
+        public void ManageInvoice()
+        {{
+            Console.WriteLine("Managing invoices..");
+        }}
+    }}
+    [/example_code]
+
+    After:
+    [example_code]
+    public class InvoiceManager
+    {{
+        public void Manage()
+        {{
+            Console.WriteLine("Managing invoices...");
+        }}
+    }}
+    [/example_code]
+
+    2. Property Naming
+    a. Public Properties:
+        - Use PascalCase (e.g., FirstName, Salary).
+        - Ensure names are concise yet descriptive. Avoid redundant prefixes that repeat the class context.
+        - Example: In a class Employee, name a property Id, not EmpId, since the class name already provides context.
+    b. Private Fields:
+        - Use a leading underscore `_` followed by camelCase (e.g., _firstName, _salary).
+    c. Logical Naming:
+        - Ensure property names clearly describe what they hold.
+        - Avoid abbreviations unless widely understood (e.g., use DateOfBirth, not DOB).
+
+    Example:
+    Before:
+    [example_code]
+    public class Employee
+    {{
+        public string empName {{ get; set; }}
+        private int emp_age;
+    }}
+    [/example_code]
+
+    After:
+    [example_code]
+    public class Employee
+    {{
+        public string Name {{get; set; }}
+        private int _age;
+
+        public void SetAge(int age)
+        {{
+            _age = age;
+        }}
+
+        public int GetAge()
+        {{
+            return _age;
+        }}
+    }}
+    [/example_code]
+
+    3. Object Naming
+    a. Use expressive names for instantiated objects to make the code self-explanatory:
+        - Example: Employee employee is more meaningful than Employee e.
+    b. Ensure names reflect their role or purpose in the code:
+        - Example: If an object calculates totals, name it totalCalculator instead of calcObj.
+
+    Example:
+    Before:
+    [example_code]
+    Employee e = new Employee();
+    e.Name = "John Doe";
+    [/example_code]
+    
+    After:
+    [example_code]
+    Employee employee = new Employee();
+    employee.Name = "John Doe";
+    [/example_code]
+
+    4. Method Naming
+    a. Use PascalCase for all method names (e.g., CalculateSalary, GetEmployeeDetails).
+    b. Ensure method names describe what the method does:
+        - Example: Use GenerateReport instead of GenReport or Process.
+    c. Avoid vague or overly generic names. Each method name should immediately convey its functionality.
+
+    Example:
+    Before:
+    [example_code]
+    public void GenReport()
+    {{
+        Console.WriteLine("Report Generated.");
+    }}
+    [/example_code]
+    
+    After:
+    [example_code]
+    public void GenerateMonthlyReport()
+    {{
+        Console.WriteLine("Monthly report has been generated.");
+    }}
+    [/example_code]
+
+    5. Method Modularity
+    a. Ensure that methods follow the Single Responsibility Principle:
+        - Break down methods that perform multiple tasks into smaller, logically focused methods.
+        - Example: A method ProcessPayroll that calculates totals, generates reports, and updates the database should be split into:
+            1. CalculatePayrollTotals
+            2. GeneratePayrollReport
+            3. UpdatePayrollDatabase
+        - Use meaningful names for all extracted methods to reflect their specific task.
+    b. Consolidate repeated logic into reusable utility methods or helper functions.
+
+    Example:
+    Before:
+    [example_code]
+    public void ProcessPayroll()
+    {{
+        Calculate();
+        Generate();
+        Update();
+    }}
+    [/example_code]
+    
+    After:
+    [example_code]
+    public void ProcessPayroll()
+    {{
+        CalculatePayrollTotals();
+        GeneratePayrollReport();
+        UpdatePayrollDatabase();
+    }}
+
+    private void CalculatePayrollTotals()
+    {{
+        Console.WriteLine("Calculating payroll totals...");
+    }}
+
+    private void GeneratePayrollReport()
+    {{
+        Console.WriteLine("Generating payroll report...");
+    }}
+
+    private void UpdatePayrollDatabase()
+    {{
+        Console.WriteLine("Updating payroll database...");
+    }}
+    [/example_code]
+
+    6. Single Responsibility Principle (SRP)
+    a. Ensure each class is responsible for only one distinct task or purpose:
+        - Example: If a class Invoice has methods for printing, calculating, and saving, split it into:
+        - Invoice (business logic, such as calculating totals).
+        - InvoicePrinter (handles formatting and output).
+        - InvoiceRepository (handles database operations).
+    b. Make classes precise and cohesive:
+        - If a class like Employee has methods like ApproveTimeOff, consider whether subclasses such as Manager or Intern would better represent specialized roles.
+        - Use inheritance to maintain logical separation of behavior.
+
+    Example:
+    Before:
+    [example_code]
+    public class Employee
+    {{
+        public void ApplyForVacation(){{/* ... */}}
+        public void ApproveTimeOff() {{ /* ... */ }}
+    }}
+
+    public static void Main(string[] args){{
+        Employee intern = new Employee();
+        Employee manager = new Manager();
+
+        intern.ApplyForVacation();
+        manager.ApproveTimeOff();
+    }}
+    [/example_code]
+    
+    After:
+    [example_code]
+    public class Employee
+    {{
+        public void ApplyForVacation(){{/* ... */}}
+    }}
+
+    public class Intern: Emplyee
+    {{
+    
+    }}
+
+    public class Manager: Employee
+    {{
+        public void ApproveTimeOff() {{ /* ... */ }}
+    }}
+
+    public static void Main(string[] args)
+    {{
+        Intern intern = new Employee();
+        Manager manager = new Manager();
+
+        intern.ApplyForVacation();
+        manager.ApproveTimeOff();
+    }}
+    [/example_code]
+
+    7. Code Clean-Up
+    a. Remove unused imports, variables, and comments to reduce clutter and improve readability.
+    b. Ensure the code is free of dead or redundant logic.
+
+    Example:
+    Before:
+    [example_code]
+    using System;
+    using System.Collections.Generic;
+    // Unused import
+    using System.Linq;
+
+    public class Employee
+    {{
+        public string Name {{ get; set; }}
+        // Commented-out code
+        // public int Age {{ get; set; }}
+    }}
+    [/example_code]
+
+    After:
+    [example_code]
+    using System;
+
+    public class Employee
+    {{
+        public string Name {{ get; set; }}
+    }}
+    [/example_code]
+
+    End Goal
+    The refined code should:
+    1. Adhere to C# naming conventions and clean code principles.
+    2. Be modular and easy to maintain, with a clear separation of concerns.
+    3. Follow the Single Responsibility Principle, ensuring each class and method has a well-defined purpose.
+    4. Be expressive, making it easy for any developer to understand the code's intent at a glance.
+
+    [C#]
+    {example["code"]}
+    [/C#]
+    """
+
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_prompt}
+    ]
+    
+    text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    no_of_input_tokens = model_inputs.input_ids.shape[1]
+    max_output_tokens =  MAX_MODEL_TOKENS - no_of_input_tokens
+        
+    
+    no_of_retries = 0
+    should_try = True
+    if max_output_tokens < no_of_input_tokens:
+        print(f'Number of input tokens is very large [{no_of_input_tokens}]. Skippping this code refinement for {example["path"]}', flush=True)
+        should_try = False
+
+    while should_try:  
+        try:
+            generated_ids = model.generate(
+                **model_inputs,
+                max_new_tokens=max_output_tokens
+            )
+            generated_ids = [
+                output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+            ]
+
+            output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+            code_pattern = r'\[refined_C#\](.*?)\[/refined_C#\]'
+            summary_pattern = r'\[code_changes\](.*?)\[/code_changes\]'
+            
+            code_matches = re.search(code_pattern, output, re.DOTALL)
+            summary_matches = re.search(summary_pattern, output, re.DOTALL)
+            
+            is_generation_okay = False
+            if code_matches and summary_matches:
+                refined_code = code_matches.group(1)
+                summary = summary_matches.group(1)
+                
+                batch_dataset.append({'code': example["code"], 'refined code': refined_code, 'summary': summary})
+                instance_number += 1
+                should_try = False
+
+            no_of_retries += 1
+            if no_of_retries == 3:
+                print(f'Could not clean the code. The final try output is{output}', flush=True)
+                should_try = False  
+        except Exception as error:
+            print(f'Error is: {error}', flush=True)
+    
+    if instance_number == BATCH_SIZE:
+        try:
+            if not is_dataset_created:
+                new_dataset = Dataset.from_generator(data_generator, gen_kwargs={"dataset": batch_dataset})
+                new_dataset.push_to_hub(DATASET_NAME)
+                is_dataset_created = True
+            else:
+                refined_code_dataset = load_dataset(DATASET_NAME_WITH_USERNAME, split="train", keep_in_memory=False)
+                new_dataset = concatenate_datasets([refined_code_dataset, Dataset.from_generator(data_generator, gen_kwargs={"dataset": batch_dataset})])
+                new_dataset.push_to_hub(DATASET_NAME)
+        except Exception as push_error:
+            print(f'Error during push to hub: {push_error}', flush=True)
+            
+
+        instance_number = 0
+        batch_dataset = []
+        batch_number += 1
+        print(f'Pushed batch number {batch_number} to hub. Current Time {datetime.datetime.now()}', flush=True)
+    
+    if batch_number == NO_OF_BATCHES:
+        break
+
+print('Dataset generation completed.')
\ No newline at end of file

data_generation/slurm.sh

+#!/usr/bin/zsh
+
+### Add basic configuration for job
+
+#SBATCH --job-name=dataset_generation
+#SBATCH --output=logs/dataset_generation_%j.log
+#SBATCH --error=logs/dataset_generation_error_%j.log
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=6
+#SBATCH --gres=gpu:1
+#SBATCH --time=20:00:00
+
+
+###------------------------------------------------------------------------------------------------------------------------------
+
+### Run the project in work directory of the cluster (configure based on need!! 
+### RWTH File System : https://help.itc.rwth-aachen.de/en/service/rhr4fjjutttf/article/da307ec2c60940b29bd42ac483fc3ea7/
+cd $HPCWORK
+cd codebud/data_generation
+###------------------------------------------------------------------------------------------------------------------------------
+
+### JOB SCRIPT RUN
+module load GCCcore/.13.2.0
+module load Python/3.11.5
+module load CUDA
+
+source ../../venvs/codebud/bin/activate
+echo $VIRTUAL_ENV
+
+python --version
+python main.py
+
+module unload CUDA
+module unload Python/3.11.5
+
+deactivate
+echo "Script ran successfully"
\ No newline at end of file

data_generation/slurm_qwen.sh

+#!/usr/bin/zsh
+
+### Add basic configuration for job
+
+#SBATCH --job-name=dataset_generation
+#SBATCH --output=logs/dataset_generation_%j.log
+#SBATCH --error=logs/dataset_generation_error_%j.log
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --time=2-15:00:00
+
+
+###------------------------------------------------------------------------------------------------------------------------------
+
+### Run the project in work directory of the cluster (configure based on need!! 
+### RWTH File System : https://help.itc.rwth-aachen.de/en/service/rhr4fjjutttf/article/da307ec2c60940b29bd42ac483fc3ea7/
+cd $HPCWORK
+cd codebud/data_generation
+###------------------------------------------------------------------------------------------------------------------------------
+
+### JOB SCRIPT RUN
+module load GCCcore/.13.2.0
+module load Python/3.11.5
+module load CUDA
+
+source ../../venvs/codebud/bin/activate
+echo $VIRTUAL_ENV
+
+python --version
+python main_qwen.py
+
+module unload CUDA
+module unload Python/3.11.5
+
+deactivate
+echo "Script ran successfully"
\ No newline at end of file

requirements.txt

+accelerate==1.1.1
+aiohappyeyeballs==2.4.3
+aiohttp==3.11.8
+aiosignal==1.3.1
+attrs==24.2.0
+certifi==2024.8.30
+charset-normalizer==3.4.0
+colorama==0.4.6
+datasets==3.1.0
+dill==0.3.8
+filelock==3.16.1
+frozenlist==1.5.0
+fsspec==2024.9.0
+huggingface-hub==0.26.3
+idna==3.10
+Jinja2==3.1.3
+MarkupSafe==2.1.5
+mpmath==1.3.0
+multidict==6.1.0
+multiprocess==0.70.16
+networkx==3.2.1
+numpy==2.1.3
+packaging==24.2
+pandas==2.2.3
+pillow==10.2.0
+propcache==0.2.0
+psutil==6.1.0
+pyarrow==18.1.0
+python-dateutil==2.9.0.post0
+pytz==2024.2
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.3
+safetensors==0.4.5
+setuptools==70.0.0
+six==1.16.0
+sympy==1.13.1
+tokenizers==0.20.3
+torch==2.5.1+cu118
+torchaudio==2.5.1+cu118
+torchvision==0.20.1+cu118
+tqdm==4.67.1
+transformers==4.46.3
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==2.2.3
+xxhash==3.5.0
+yarl==1.18.0