Update README.md

README.md

@@ -1,4 +1,70 @@
 ---
 license: apache-2.0
 ---
-this is a demo how to pretrain a mistral architecture model by SFT Trainer within tens of lines Python code.
+This is a demo of how to pretrain a mistral architecture model by SFT Trainer ,and it needs only 70 lines Python code.
+
+```
+import torch
+from transformers import TrainingArguments, MistralForCausalLM, MistralModel, MistralConfig, AutoTokenizer
+from datasets import load_dataset
+from trl import SFTTrainer
+
+configuration = MistralConfig(vocab_size=32000,
+        hidden_size=2048,
+        intermediate_size=7168,
+        num_hidden_layers=24,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        pad_token_id=2,
+        bos_token_id=1,
+        eos_token_id=2)
+
+model = MistralForCausalLM(configuration)
+#model = MistralForCausalLM.from_pretrained("./6B_code_outputs/checkpoint-10000")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2", local_files_only=False)
+tokenizer.pad_token = tokenizer.eos_token
+
+dataset = load_dataset('HuggingFaceTB/cosmopedia-20k', split="train")
+#dataset = load_dataset('Elriggs/openwebtext-100k', split="train")
+dataset = dataset.shuffle(seed=42)
+print(f'Number of prompts: {len(dataset)}')
+print(f'Column names are: {dataset.column_names}')
+
+def create_prompt_formats(sample):
+    """
+    Format various fields of the sample ('instruction', 'context', 'response')
+    Then concatenate them using two newline characters
+    :param sample: Sample dictionnary
+    """
+    output_texts = []
+    for i in range(len(sample['text'])):
+      formatted_prompt = sample['text'][i]
+      output_texts.append(formatted_prompt)
+    #print(output_texts)
+    return output_texts
+
+
+trainer = SFTTrainer(
+    model,
+    train_dataset=dataset,
+    tokenizer = tokenizer,
+    max_seq_length=2048,
+    formatting_func=create_prompt_formats,
+    args=TrainingArguments(
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=1,
+            warmup_steps=2,
+            max_steps=10000,
+            learning_rate=1e-4,
+            logging_steps=1,
+            output_dir="6B_outputs", overwrite_output_dir=True,save_steps=1000,
+            optim="paged_adamw_32bit",report_to="none"
+        )
+)
+trainer.train()
+trainer.model.save_pretrained("6B-final", dtype=torch.float32)
+trainer.tokenizer.save_pretrained("6B-final")
+
+```