Update README.md

README.md

@@ -71,19 +71,39 @@ Here is how you can use this model:
 ```python
 from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
+# Define the base model and the adapter model
 base_model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
 adapter_model = "AmirMohseni/Llama-3.1-8B-Instruct-Persian-finetuned-sft"
 
-model = AutoModelForCausalLM.from_pretrained(base_model)
+# Load the base model and apply the adapter model using PEFT
+model = AutoModelForCausalLM.from_pretrained(base_model, device_map={"": 0})
 model = PeftModel.from_pretrained(model, adapter_model)
 
+# Check if CUDA is available, otherwise use CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+# Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model)
 
+# Add a new pad token if necessary
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})  # Adding a distinct pad token
+
 # Example usage
-prompt = "راه‌های تقویت حافظه چیست؟"
-inputs = tokenizer(prompt, return_tensors="pt")
-outputs = model.generate(**inputs)
+input_text = "چطوری میتونم به اطلاعات درباره ی سهام شرکت های آمریکایی دست پیدا کنم؟"
+
+# Tokenize the input and get both input IDs and attention mask
+inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
+input_ids = inputs['input_ids'].to(device)
+attention_mask = inputs['attention_mask'].to(device)
+
+# Generate text
+outputs = model.generate(input_ids, attention_mask=attention_mask, max_length=512, pad_token_id=tokenizer.pad_token_id)
+
+# Decode and print the output
 response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 print(response)
 ```