Update README.md

README.md

@@ -84,7 +84,7 @@ repo_id = "jordiclive/lora-llama-33B-alpaca_gpt4-dolly_15k-vicuna-r64"
 base_model = "decapoda-research/llama-30b-hf"
 
 # Model Loading
-def transfer_embeddings(model, embed_path, tokenizer):
+def add_embeddings(model, embed_path, tokenizer):
     old_embeddings = model.get_input_embeddings()
     old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
     new_embeddings = torch.nn.Embedding(old_num_tokens, old_embedding_dim)
@@ -93,16 +93,17 @@ def transfer_embeddings(model, embed_path, tokenizer):
     embed_weights = torch.load(embed_path, map_location=old_embeddings.weight.device)
     vocab_size = tokenizer.vocab_size
     new_embeddings.weight.data[:vocab_size, :] = old_embeddings.weight.data[:vocab_size, :]
-    new_embeddings.weight.data[vocab_size : vocab_size + embed_weights.shape[0], :] = embed_weights.weight.data.to(
+    new_embeddings.weight.data[vocab_size : vocab_size + embed_weights.shape[0], :] = embed_weights.to(
         new_embeddings.weight.dtype
     ).to(new_embeddings.weight.device)
     model.set_input_embeddings(new_embeddings)
     model.tie_weights()
 
 
+
 def load_peft_model(model, peft_model_path, tokenizer):
     embed_weights = hf_hub_download(peft_model_path, "extra_embeddings.pt")
-    model.resize_token_embeddings(tokenizer.vocab_size + embed_weights.shape[0])
+    model.resize_token_embeddings(tokenizer.vocab_size + torch.load(embed_weights).shape[0])
     model.config.eos_token_id = tokenizer.eos_token_id
     model.config.bos_token_id = tokenizer.bos_token_id
     model.config.pad_token_id = tokenizer.pad_token_id
@@ -112,20 +113,22 @@ def load_peft_model(model, peft_model_path, tokenizer):
         torch_dtype=model.dtype,
     )
     model.eos_token_id = tokenizer.eos_token_id
-    transfer_embeddings(model, Path(peft_model_path).joinpath("extra_embeddings.pt"), tokenizer)
+    add_embeddings(model, embed_weights, tokenizer)
     return model
 
 
 tokenizer = transformers.AutoTokenizer.from_pretrained(repo_id)
 
 model = transformers.AutoModelForCausalLM.from_pretrained(
-    base_model, torch_dtype=dtype, trust_remote_code=True, cache_dir="/mnt/data/jordiclive/data_cache"
+    base_model, torch_dtype=dtype, trust_remote_code=True,
 )
 model = load_peft_model(model, repo_id, tokenizer)
 
 
 # device  configuration
 model = model.to(device)
+if dtype == torch.float16:
+    model = model.half()
 
 
 # Choose Generation parameters
@@ -164,6 +167,4 @@ def generate(prompt, generation_config=generation_config, max_new_tokens=2048, d
 generate("What is a meme, and what's the history behind this word?")
 generate("What's the Earth total population")
 generate("Write a story about future of AI development")
-
-
 ```