





'''simple docstring'''





import re
import tempfile
from pathlib import Path

import pytest
import yaml

from datasets.utils.readme import ReadMe


# @pytest.fixture
# def example_yaml_structure():

UpperCamelCase__   : Optional[Any]					  = yaml.safe_load(
    """\
name: \"\"
allow_empty: false
allow_empty_text: true
subsections:
  - name: \"Dataset Card for X\" # First-level markdown heading
    allow_empty: false
    allow_empty_text: true
    subsections:
      - name: \"Table of Contents\"
        allow_empty: false
        allow_empty_text: false
        subsections: null
      - name: \"Dataset Description\"
        allow_empty: false
        allow_empty_text: false
        subsections:
          - name: \"Dataset Summary\"
            allow_empty: false
            allow_empty_text: false
            subsections: null
          - name: \"Supported Tasks and Leaderboards\"
            allow_empty: true
            allow_empty_text: true
            subsections: null
          - name: Languages
            allow_empty: false
            allow_empty_text: true
            subsections: null
"""
)


UpperCamelCase__   : Any					  = {
    """name""": """root""",
    """text""": """""",
    """is_empty_text""": True,
    """subsections""": [
        {
            """name""": """Dataset Card for My Dataset""",
            """text""": """""",
            """is_empty_text""": True,
            """subsections""": [
                {"""name""": """Table of Contents""", """text""": """Some text here.""", """is_empty_text""": False, """subsections""": []},
                {
                    """name""": """Dataset Description""",
                    """text""": """Some text here.""",
                    """is_empty_text""": False,
                    """subsections""": [
                        {
                            """name""": """Dataset Summary""",
                            """text""": """Some text here.""",
                            """is_empty_text""": False,
                            """subsections""": [],
                        },
                        {
                            """name""": """Supported Tasks and Leaderboards""",
                            """text""": """""",
                            """is_empty_text""": True,
                            """subsections""": [],
                        },
                        {"""name""": """Languages""", """text""": """Language Text""", """is_empty_text""": False, """subsections""": []},
                    ],
                },
            ],
        }
    ],
}


UpperCamelCase__   : Optional[Any]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""


UpperCamelCase__   : Union[str, Any]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
#### Extra Ignored Subsection
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : str					  = {
    """name""": """root""",
    """text""": """""",
    """is_empty_text""": True,
    """subsections""": [
        {
            """name""": """Dataset Card for My Dataset""",
            """text""": """""",
            """is_empty_text""": True,
            """subsections""": [
                {"""name""": """Table of Contents""", """text""": """Some text here.""", """is_empty_text""": False, """subsections""": []},
                {
                    """name""": """Dataset Description""",
                    """text""": """Some text here.""",
                    """is_empty_text""": False,
                    """subsections""": [
                        {
                            """name""": """Dataset Summary""",
                            """text""": """Some text here.""",
                            """is_empty_text""": False,
                            """subsections""": [
                                {
                                    """name""": """Extra Ignored Subsection""",
                                    """text""": """""",
                                    """is_empty_text""": True,
                                    """subsections""": [],
                                }
                            ],
                        },
                        {
                            """name""": """Supported Tasks and Leaderboards""",
                            """text""": """""",
                            """is_empty_text""": True,
                            """subsections""": [],
                        },
                        {"""name""": """Languages""", """text""": """Language Text""", """is_empty_text""": False, """subsections""": []},
                    ],
                },
            ],
        }
    ],
}

UpperCamelCase__   : Tuple					  = """\
---
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : List[Any]					  = (
    """The following issues were found for the README at `{path}`:\n-\tEmpty YAML markers are present in the README."""
)

UpperCamelCase__   : Dict					  = """\
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : List[str]					  = (
    """The following issues were found for the README at `{path}`:\n-\tNo YAML markers are present in the README."""
)

UpperCamelCase__   : Dict					  = """\
---
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : Optional[Any]					  = """The following issues were found for the README at `{path}`:\n-\tOnly the start of YAML tags present in the README."""

UpperCamelCase__   : List[str]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
### Supported Tasks and Leaderboards
### Languages
Language Text
"""
UpperCamelCase__   : Any					  = """The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Summary` but it is empty.\n-\tExpected some text in section `Dataset Summary` but it is empty (text in subsections are ignored)."""


UpperCamelCase__   : str					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
"""
UpperCamelCase__   : Dict					  = """The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Card for My Dataset` but it is empty.\n-\tSection `Dataset Card for My Dataset` expected the following subsections: `Table of Contents`, `Dataset Description`. Found 'None'."""

UpperCamelCase__   : Union[str, Any]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Languages
Language Text
"""

UpperCamelCase__   : int					  = """The following issues were found for the README at `{path}`:\n-\tSection `Dataset Description` is missing subsection: `Supported Tasks and Leaderboards`."""


UpperCamelCase__   : List[str]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
"""

UpperCamelCase__   : Optional[int]					  = """The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Languages` but it is empty."""

UpperCamelCase__   : List[Any]					  = """\
---
language:
- zh
- en
---

## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""
UpperCamelCase__   : int					  = """The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README."""

UpperCamelCase__   : Any					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
# Dataset Card My Dataset
"""

UpperCamelCase__   : List[Any]					  = """The following issues were found for the README at `{path}`:\n-\tThe README has several first-level headings: `Dataset Card for My Dataset`, `Dataset Card My Dataset`. Only one heading is expected. Skipping further validation for this README."""

UpperCamelCase__   : Optional[int]					  = """\
---
language:
- zh
- en
---

# Dataset Card My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : Tuple					  = """The following issues were found for the README at `{path}`:\n-\tNo first-level heading starting with `Dataset Card for` found in README. Skipping further validation for this README."""

UpperCamelCase__   : Optional[int]					  = """"""

UpperCamelCase__   : List[str]					  = """The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.\n-\tNo YAML markers are present in the README."""

UpperCamelCase__   : Optional[Any]					  = """\
---
language:
- zh
- en
---

# Dataset Card for My Dataset
# Dataset Card for My Dataset
## Table of Contents
Some text here.
## Dataset Description
Some text here.
### Dataset Summary
Some text here.
### Supported Tasks and Leaderboards
### Languages
Language Text
"""

UpperCamelCase__   : str					  = """The following issues were found while parsing the README at `{path}`:\n-\tMultiple sections with the same heading `Dataset Card for My Dataset` have been found. Please keep only one of these sections."""




@pytest.mark.parametrize(
    """readme_md, expected_dict"""				,      [
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,      )
def    A_(			A				,      A   ):
	assert ReadMe.from_string(A				,      A   ).to_dict() == expected_dict




@pytest.mark.parametrize(
    """readme_md, expected_error"""				,      [
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,      )
def    A_(			A				,      A   ):
	with pytest.raises(A				,      match=re.escape(expected_error.format(path="""root"""   )   )   ):
		UpperCAmelCase_								=     ReadMe.from_string(A				,      A   )
		readme.validate()




@pytest.mark.parametrize(
    """readme_md, expected_error"""				,      [
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,      )
def    A_(			A				,      A   ):
	with pytest.raises(A				,      match=re.escape(expected_error.format(path="""root"""   )   )   ):
		ReadMe.from_string(A				,      A   )




@pytest.mark.parametrize(
    """readme_md,"""				,      [
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,      )
def    A_(			A   ):
	ReadMe.from_string(A				,      A				,      suppress_parsing_errors=A   )




@pytest.mark.parametrize(
    """readme_md, expected_dict"""				,      [
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,      )
def    A_(			A				,      A   ):
	with tempfile.TemporaryDirectory() as tmp_dir:
		UpperCAmelCase_								=     Path(A   ) / """README.md"""
		with open(A				,      """w+"""   ) as readme_file:
			readme_file.write(A   )
		UpperCAmelCase_								=     ReadMe.from_readme(A				,      A   ).to_dict()
		assert out["name"] == path
		assert out["text"] == ""
		assert out["is_empty_text"]
		assert out["subsections"] == expected_dict["subsections"]




@pytest.mark.parametrize(
    """readme_md, expected_error"""				,      [
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,      )
def    A_(			A				,      A   ):
	with tempfile.TemporaryDirectory() as tmp_dir:
		UpperCAmelCase_								=     Path(A   ) / """README.md"""
		with open(A				,      """w+"""   ) as readme_file:
			readme_file.write(A   )
		UpperCAmelCase_								=     expected_error.format(path=A   )
		with pytest.raises(A				,      match=re.escape(A   )   ):
			UpperCAmelCase_								=     ReadMe.from_readme(A				,      A   )
			readme.validate()




@pytest.mark.parametrize(
    """readme_md, expected_error"""				,      [
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,      )
def    A_(			A				,      A   ):
	with tempfile.TemporaryDirectory() as tmp_dir:
		UpperCAmelCase_								=     Path(A   ) / """README.md"""
		with open(A				,      """w+"""   ) as readme_file:
			readme_file.write(A   )
		UpperCAmelCase_								=     expected_error.format(path=A   )
		with pytest.raises(A				,      match=re.escape(A   )   ):
			ReadMe.from_readme(A				,      A   )




@pytest.mark.parametrize(
    """readme_md,"""				,      [
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,      )
def    A_(			A   ):
	with tempfile.TemporaryDirectory() as tmp_dir:
		UpperCAmelCase_								=     Path(A   ) / """README.md"""
		with open(A				,      """w+"""   ) as readme_file:
			readme_file.write(A   )
		ReadMe.from_readme(A				,      A				,      suppress_parsing_errors=A   )









import argparse
from copy import deepcopy

import numpy as np
from datasets import ClassLabel, DatasetDict, load_dataset
from evaluate import load

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    set_seed,
)




def    A_(			):
	UpperCAmelCase_								=     argparse.ArgumentParser()
	parser.add_argument("""--model_ckpt"""				,      type=A				,      default="""microsoft/unixcoder-base-nine"""   )
	parser.add_argument("""--num_epochs"""				,      type=A				,      default=5   )
	parser.add_argument("""--batch_size"""				,      type=A				,      default=6   )
	parser.add_argument("""--gradient_accumulation_steps"""				,      type=A				,      default=1   )
	parser.add_argument("""--freeze"""				,      type=A				,      default=A   )
	parser.add_argument("""--learning_rate"""				,      type=A				,      default=5E-4   )
	parser.add_argument("""--seed"""				,      type=A				,      default=0   )
	parser.add_argument("""--lr_scheduler_type"""				,      type=A				,      default="""cosine"""   )
	parser.add_argument("""--num_warmup_steps"""				,      type=A				,      default=10   )
	parser.add_argument("""--weight_decay"""				,      type=A				,      default=0.01   )
	parser.add_argument("""--output_dir"""				,      type=A				,      default="""./results"""   )
	return parser.parse_args()


UpperCamelCase__   : Any					  = load("""accuracy""")




def    A_(			A   ):
	UpperCAmelCase_						,	UpperCAmelCase_								=     eval_pred
	UpperCAmelCase_								=     np.argmax(A				,      axis=1   )
	return metric.compute(predictions=A				,      references=A   )







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					def __init__(						self    :  str ,					__lowercase    :  Any		):






						'''simple docstring'''







						super().__init__()
						UpperCAmelCase_								=     trainer





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple ,					__lowercase    :  int ,					__lowercase    :  Tuple ,					__lowercase    :  Dict ,					**__lowercase    :  Dict		):






						'''simple docstring'''







						if control.should_evaluate:
							UpperCAmelCase_								=     deepcopy(__lowercase		)
							self._trainer.evaluate(eval_dataset=self._trainer.train_dataset ,					metric_key_prefix="""train"""		)
							return control_copy




def    A_(			):
	UpperCAmelCase_								=     get_args()
	set_seed(args.seed   )

	UpperCAmelCase_								=     load_dataset("""codeparrot/codecomplex"""				,      split="""train"""   )
	UpperCAmelCase_								=     dataset.train_test_split(test_size=0.2   )
	UpperCAmelCase_								=     train_test["""test"""].train_test_split(test_size=0.5   )
	UpperCAmelCase_								=     DatasetDict(
	    {
	        """train""": train_test["""train"""],
	        """test""": test_validation["""train"""],
	        """valid""": test_validation["""test"""],
	    }   )

	print("""Loading tokenizer and model"""   )
	UpperCAmelCase_								=     AutoTokenizer.from_pretrained(args.model_ckpt   )
	UpperCAmelCase_								=     tokenizer.eos_token
	UpperCAmelCase_								=     AutoModelForSequenceClassification.from_pretrained(args.model_ckpt				,      num_labels=7   )
	UpperCAmelCase_								=     model.config.eos_token_id

	if args.freeze:
		for param in model.roberta.parameters():
			UpperCAmelCase_								=     False

	UpperCAmelCase_								=     ClassLabel(num_classes=7				,      names=list(set(train_test_validation["""train"""]["""complexity"""]   )   )   )

	def tokenize(A   ):
		UpperCAmelCase_								=     tokenizer(example["""src"""]				,      truncation=A				,      max_length=1024   )
		UpperCAmelCase_								=     labels.straint(example["""complexity"""]   )
		return {
		    "input_ids": inputs["input_ids"],
		    "attention_mask": inputs["attention_mask"],
		    "label": label,
		}

	UpperCAmelCase_								=     train_test_validation.map(
	    A				,      batched=A				,      remove_columns=train_test_validation["""train"""].column_names				,      )
	UpperCAmelCase_								=     DataCollatorWithPadding(tokenizer=A   )

	UpperCAmelCase_								=     TrainingArguments(
	    output_dir=args.output_dir				,      learning_rate=args.learning_rate				,      lr_scheduler_type=args.lr_scheduler_type				,      evaluation_strategy="""epoch"""				,      save_strategy="""epoch"""				,      logging_strategy="""epoch"""				,      per_device_train_batch_size=args.batch_size				,      per_device_eval_batch_size=args.batch_size				,      num_train_epochs=args.num_epochs				,      gradient_accumulation_steps=args.gradient_accumulation_steps				,      weight_decay=0.01				,      metric_for_best_model="""accuracy"""				,      run_name="""complexity-java"""				,      report_to="""wandb"""				,      )

	UpperCAmelCase_								=     Trainer(
	    model=A				,      args=A				,      train_dataset=tokenized_datasets["""train"""]				,      eval_dataset=tokenized_datasets["""valid"""]				,      tokenizer=A				,      data_collator=A				,      compute_metrics=A				,      )

	print("""Training..."""   )
	trainer.add_callback(CustomCallback(A   )   )
	trainer.train()


if __name__ == "__main__":
				main()


