Update README.md

README.md

@@ -34,15 +34,18 @@ This model is trained based on the scripts provided in https://github.com/tloen/
 ```Python
 # adapted from https://github.com/tloen/alpaca-lora/blob/main/generate.py
 
+import torch
 from peft import PeftModel
+from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
+
 
 if torch.cuda.is_available():
     device = "cuda"
 else:
     device = "cpu"
 
-tokenizer = LLaMATokenizer.from_pretrained("decapoda-research/llama-7b-hf")
-base_model = LLaMAForCausalLM.from_pretrained(
+tokenizer = LlaMATokenizer.from_pretrained("decapoda-research/llama-7b-hf")
+base_model = LlaMAForCausalLM.from_pretrained(
     "decapoda-research/llama-7b-hf",
     load_in_8bit=True,
     device_map="auto",