Update README.md

README.md

@@ -38,13 +38,23 @@ This model can be easily loaded using the `AutoModelForCausalLM` functionality:
 
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
-tokenizer = AutoTokenizer.from_pretrained("ammmarnasr/codegen-350M-mono-swift")
-model = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono")
+from peft import PeftConfig, PeftModel
 
-text = "def hello_world():"
-input_ids = tokenizer(text, return_tensors="pt").input_ids
+model_name = "ammarnasr/codegen-350M-mono-swift"
+peft_config = PeftConfig.from_pretrained(model_name)
 
-generated_ids = model.generate(input_ids, max_length=128)
+tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
+
+model = AutoModelForCausalLM.from_pretrained(peft_config.base_model_name_or_path)
+model = PeftModel.from_pretrained(model, model_name)
+
+model.print_trainable_parameters()
+
+text = "func hello_world() {"
+
+input_ids = tokenizer.encode(text, return_tensors="pt")
+generated_ids = model.generate(input_ids=input_ids, max_length=100)
+print('Generated: \n')
 print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
 ```