Update README.md

README.md

@@ -74,7 +74,7 @@ repo_id = "jordiclive/falcon_lora_40b_ckpt_500_oasst_1"
 base_model = "tiiuae/falcon-40b"
 
 # 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)
@@ -83,16 +83,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
@@ -102,7 +103,7 @@ def load_peft_model(model, peft_model_path, tokenizer):
         torch_dtype=model.dtype,
     )
     model.eos_token_id = tokenizer.eos_token_id
-    transfer_embeddings(model, peft_model_path.joinpath("extra_embeddings.pt"), tokenizer)
+    add_embeddings(model, embed_weights, tokenizer)
     return model
 
 
@@ -116,6 +117,8 @@ model = load_peft_model(model, repo_id, tokenizer)
 
 # device  configuration
 model = model.to(device)
+if dtype == torch.float16:
+    model = model.half()
 
 
 # Choose Generation parameters
@@ -155,4 +158,5 @@ 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")
 
+
 ```