im making fillin middle dataset 
User
Here is the code make necessary changes:

import torch from torch.utils.data import DataLoader from datasets import Dataset, DatasetDict import pandas as pd import random from transformers import AutoModelForCausalLM, AutoTokenizer, DataCollatorForLanguageModeling

model = "Qwen/Qwen1.5-1.8B" tokenizer = AutoTokenizer.from_pretrained(model) model = AutoModelForCausalLM.from_pretrained(model) tokenizer.pad_token = tokenizer.eos_token

from datasets import load_dataset prefix_token = "<fim_prefix>" middle_token = "<fim_middle>" suffix_token = "<fim_suffix>" eot_token = "" context_length = 512 chunk_length = context_length max_len = context_length ds_train = load_dataset("huggingface-course/codeparrot-ds-train", split="train") ds_valid = load_dataset("huggingface-course/codeparrot-ds-valid", split="validation") raw_datasets = DatasetDict( { "train": ds_train.shuffle().select(range(50000)), "valid": ds_valid.shuffle().select(range(500)), } )

def split_document(doc, fim_rate=0.5):

if random.random() < fim_rate:

length = len(doc)

prefix_len = random.randint(0,length)

suffix_len = random.randint(0, length - prefix_len)

middle_len = length - prefix_len - suffix_len

prefix = doc[:prefix_len]

middle = doc[prefix_len:prefix_len+middle_len]

suffix = doc[prefix_len+middle_len:]

return prefix, middle, suffix

else:

return doc, None, None

def format_psm(prefix, middle, suffix, tokenizer): fromatted_examle = f"{prefix_token}{prefix}{suffix_token}{suffix}{middle_token}{middle}{eot_token}" return formatted_examle

def format_spm(prefix, middle, suffix, tokenizer): fromatted_examle = f"{prefix_token}{suffix_token}{suffix}{middle_token}{prefix}{middle}{eot_token}" return formatted_examle

def format_nofim(doc, tokenizer): formatted_examle = f"{doc}" return formatted_examle

def apply_fim_transformation(chunk_list, p_psm=0.5): transformed_docs = [] for chunk in chunk_list: prefix, middle, suffix = split_document(chunk) if middle is not None: if random.random() < p_psm: transformed_doc = format_psm(prefix, middle, suffix, tokenizer) else: transformed_doc = format_spm(prefix, middle, suffix, tokenizer) transformed_docs.append(transformed_doc) else: transformed_doc = format_nofim(chunk, tokenizer) transformed_docs.append(transformed_doc) return transformed_docs

def join_transformed_chunk_docs(transformed_chunk_docs): merged_docs = transformed_chunk_docs[0] if len(transformed_chunk_docs) == 1: return merged_docs for i in range(2, len(transformed_chunk_docs)): merged_docs = merged_docs + eot_token + transformed_chunk_docs[i] return merged_docs

def apply_context_level_fim(chunk):

chunk_docs = chunk.split(eot_token)

transformed_chunk_docs = apply_fim_transformation(chunk_docs)

joined_transformed_chunk_docs = join_transformed_chunk_docs(transformed_chunk_docs)

# Return the transformed text directly

return joined_transformed_chunk_docs

class FimDataset(Dataset): def __init__(self, data): self._data = data def __len__(self): return len(self._data) def __repr__(self): return f"Dataset(num_rows={len(self)})" def __getitem__(self, index): chunk = self._data[index]["content_chunk"] chunk = str(chunk) return apply_context_level_fim(chunk)

train_dataset = FimDataset(chunk_ds["train"])

valid_dataset = FimDataset(chunk_ds["valid"])