Update README.md

README.md

@@ -51,6 +51,9 @@ from transformers import BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
 import torch
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, pipeline
+import logging
+# Suppress all warnings
+logging.getLogger("transformers").setLevel(logging.CRITICAL) #weird warning when using model for inference
 
 # Check if CUDA is available
 if torch.cuda.is_available():
@@ -70,6 +73,10 @@ base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
 model = PeftModel.from_pretrained(base_model, "smartinez1/Llama-3.1-8B-FINLLM")
 # Load the tokenizer associated with the base model
 tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
+# Define the unique padding token for fine-tuning
+custom_pad_token = "<|finetune_right_pad_id|>"
+tokenizer.add_special_tokens({'pad_token': custom_pad_token})
+pad_token_id = tokenizer.pad_token_id
 
 # Set up the text generation pipeline with the PEFT model, specifying the device
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)