Update README.md

README.md

@@ -40,13 +40,23 @@ paraphraser.paraphrase("Hey, can yuo hepl me cancel my last order?", tone="witty
 
 OR use directly with transformers
 
-```
+```python
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
 
-
+model_id = "llm-toys/RedPajama-INCITE-Base-3B-v1-paraphrase-tone"
 DEVICE = "cuda"
 EOC_FORMAT = "\n\n### END"
+PARAPHRASE_PREDICT_FORMAT = (
+    "### Instruction:\nGenerate a paraphrase for the following Input sentence.\n\n"
+    "### Input:\n{input_text}\n\n### Response:\n"
+)
+TONE_CHANGE_PREDICT_FORMAT = (
+    "### Instruction:\nChange the tone of the following Input sentence to {tone}.\n\n"
+    "### Input:\n{input_text}\n\n### Response:\n"
+)
 
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModel.from_pretrained(mode_id).to(DEVICE)
 
 class StoppingCriteriaSub(StoppingCriteria):
   """Helps in stopping the generation when a certain sequence of tokens is generated."""
@@ -63,7 +73,6 @@ stopping_criteria = StoppingCriteriaList(
   [StoppingCriteriaSub(stops=tokenizer(EOC_FORMAT)["input_ids"])]
 )
 
-
 def predict(input_text: str) -> str:
   tokenized = tokenizer(
       input_text,
@@ -86,9 +95,14 @@ def predict(input_text: str) -> str:
           stopping_criteria=self.stopping_criteria,
       )
 
-  out_texts = [self.tokenizer.decode(o, skip_special_tokens=True) for o in out]
-  for o in out_texts:
-    print(o)
+  out_texts = [tokenizer.decode(o, skip_special_tokens=True) for o in out]
+  return out_texts
+
+print("Paraphrasing:")
+print(predict(PARAPHRASE_PREDICT_FORMAT.format(input_text="If you have any further questions, feel free to ask.")))
+
+print("Tone change:")
+print(predict(TONE_CHANGE_PREDICT_FORMAT.format(input_text="If you have any further questions, feel free to ask.", tone="professional")))
 ```
 
 ## Sample training data