Update README.md

README.md

@@ -54,6 +54,12 @@ The code below performs the following steps:
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
+# Load the base pre-trained model
+base_model_id = "EleutherAI/gpt-neox-20b"
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+model = AutoModelForCausalLM.from_pretrained(base_model_id)
+
+# Fine-tuning model
 model_id = "hipnologo/GPT-Neox-20b-QLoRA-FineTune-english_quotes_dataset"
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -62,7 +68,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+# Load the fine-tuned model
 model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, device_map={"":0})
 
 text = "Twenty years from now"