Update README.md

README.md

@@ -106,12 +106,13 @@ We recommend users of this model to develop guardrails and to take appropriate p
 
 This requires a GPU with at least 27GB memory.
 
+### First, Load the Model
+
 ```python
 import torch
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# load the model
 peft_model_id = "dfurman/falcon-40b-chat-oasst1"
 config = PeftConfig.from_pretrained(peft_model_id)
 
@@ -134,8 +135,11 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 tokenizer.pad_token = tokenizer.eos_token
 
 model = PeftModel.from_pretrained(model, peft_model_id)
+```
 
-# run the model
+### Next, Run the Model
+
+```python
 prompt = """<human>: My name is Daniel. Write a short email to my closest friends inviting them to come to my home on Friday for a dinner party, I will make the food but tell them to BYOB.
 <bot>:"""
 
@@ -149,13 +153,17 @@ batch = batch.to('cuda:0')
 
 with torch.cuda.amp.autocast():
     output_tokens = model.generate(
-        input_ids = batch.input_ids, 
+        inputs=batch.input_ids, 
         max_new_tokens=200,
-        temperature=0.7,
-        top_p=0.7,
+        do_sample=False,
+        use_cache=True,
+        temperature=1.0,
+        top_k=50,
+        top_p=1.0,
         num_return_sequences=1,
         pad_token_id=tokenizer.eos_token_id,
         eos_token_id=tokenizer.eos_token_id,
+        bos_token_id=tokenizer.eos_token_id,
     )
 
 generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)