Update app.py

app.py

@@ -1,63 +1,116 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed, pipeline
-
-
-title = "Code Generator"
-description = "This is a space to convert english text to Python code using with [codeparrot-small-text-to-code](https://huggingface.co/codeparrot/codeparrot-small-text-to-code),\
-            a code generation model for Python finetuned on [github-jupyter-text](https://huggingface.co/datasets/codeparrot/github-jupyter-text) a dataset of doctrings\
-            and their Python code extracted from Jupyter notebooks."
-example = [
-    ["Utility function to compute the accuracy of predictions using metric from sklearn", 65, 0.6, 42],
-    ["Let's implement a function that computes the size of a file called filepath", 60, 0.6, 42],
-    ["Let's implement bubble sort in a helper function:", 87, 0.6, 42],
-    ]
-
-# change model to the finetuned one
-tokenizer = AutoTokenizer.from_pretrained("codeparrot/codeparrot-small-text-to-code")
-model = AutoModelForCausalLM.from_pretrained("codeparrot/codeparrot-small-text-to-code")
-
-def make_doctring(gen_prompt):
-    return "\"\"\"\n" + gen_prompt + "\n\"\"\"\n\n"
-
-def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
-    set_seed(seed)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    prompt = make_doctring(gen_prompt)
-    generated_text = pipe(prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
-    return generated_text
-
-
-iface = gr.Interface(
-    fn=code_generation, 
-    inputs=[
-        gr.Code(lines=10, language="python", label="English instructions"),
-        gr.inputs.Slider(
-            minimum=8,
-            maximum=256,
-            step=1,
-            default=8,
-            label="Number of tokens to generate",
-        ),
-        gr.inputs.Slider(
-            minimum=0,
-            maximum=2.5,
-            step=0.1,
-            default=0.6,
-            label="Temperature",
-        ),
-        gr.inputs.Slider(
-            minimum=0,
-            maximum=1000,
-            step=1,
-            default=42,
-            label="Random seed to use for the generation"
-        )
-    ],
-    outputs=gr.Code(label="Predicted Python code", language="python", lines=10),
-    examples=example,
-    layout="horizontal",
-    theme="peach",
-    description=description,
-    title=title
+from transformers import AutoTokenizer, DataCollatorForLanguageModeling, GPT2LMHeadModel, AutoConfig
+from datasets import load_dataset, DatasetDict
+
+# 加载数据集
+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)),
+        "train": ds_train.shuffle().select(range(10000)),
+      # 验证集
+        # "valid": ds_valid,  # .shuffle().select(range(500))
+        "valid": ds_valid.shuffle().select(range(500))
+    }
+)
+
+context_length = 128
+tokenizer = AutoTokenizer.from_pretrained("huggingface-course/code-search-net-tokenizer")
+
+outputs = tokenizer(
+    # 从训练集数据集中选择前两个样本的"content"字段
+    raw_datasets["train"][:2]["content"],
+    # 截断操作，如果文本长度超过max_length，则截断到指定的最大长度
+    truncation=True,
+    # 128
+    max_length=context_length,
+    # 表示如果文本长度超过了max_length，则返回超出部分的标记
+    return_overflowing_tokens=True,
+    # 表示返回每个样本处理后的标记序列的长度
+    return_length=True,
+)
+
+print(f"Input IDs length: {len(outputs['input_ids'])}")
+print(f"Input chunk lengths: {(outputs['length'])}")
+print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")
+
+def tokenize(element):
+    outputs = tokenizer(
+        element["content"],
+        truncation=True,
+        max_length=context_length,
+        return_overflowing_tokens=True,
+        return_length=True,
+    )
+    input_batch = []
+    for length, input_ids in zip(outputs["length"], outputs["input_ids"]):
+        if length == context_length:
+            input_batch.append(input_ids)
+    return {"input_ids": input_batch}
+
+
+tokenized_datasets = raw_datasets.map(
+    tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
+)
+
+print(tokenized_datasets)
+
+# 创建一个GPT-2语言模型的配置（config）对象
+config = AutoConfig.from_pretrained(
+    "gpt2",
+    vocab_size=len(tokenizer),
+    n_ctx=context_length,
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
 )
-iface.launch()
+# 初始化模型
+model = GPT2LMHeadModel(config)
+# 参数数量
+model_size = sum(t.numel() for t in model.parameters())
+
+print(f"GPT-2 size: {model_size/1000**2:.1f}M parameters")
+
+# 将分词器（tokenizer）的填充标记（pad token）设置为结束标记（eos token）
+# 这将确保在数据收集过程中，将结束标记用作填充标记，以便对不同长度的序列进行批处理。
+tokenizer.pad_token = tokenizer.eos_token
+# 用于语言建模任务的数据收集器对象
+data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+
+out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
+for key in out:
+    print(f"{key} shape: {out[key].shape}")
+
+from transformers import Trainer, TrainingArguments
+
+args = TrainingArguments(
+    output_dir="codeparrot-ds",
+    per_device_train_batch_size=32,
+    per_device_eval_batch_size=32,
+    evaluation_strategy="steps",
+    eval_steps=5_000,
+    logging_steps=5_000,
+    gradient_accumulation_steps=8,
+    num_train_epochs=1,
+    weight_decay=0.1,
+    warmup_steps=1_000,
+    lr_scheduler_type="cosine",
+    learning_rate=5e-4,
+    save_steps=5_000,
+    fp16=False,
+    push_to_hub=False,
+)
+
+trainer = Trainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=args,
+    data_collator=data_collator,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["valid"],
+)
+
+print(trainer)
+trainer.train()