Update README.md

README.md

@@ -38,6 +38,59 @@ paraphraser.paraphrase("Hey, can yuo hepl me cancel my last order?", tone="witty
 # "Hey, I need your help with my last order. Can you wave your magic wand and make it disappear?"
 ```
 
+OR use directly with transformers
+
+```
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
+
+
+DEVICE = "cuda"
+EOC_FORMAT = "\n\n### END"
+
+
+class StoppingCriteriaSub(StoppingCriteria):
+  """Helps in stopping the generation when a certain sequence of tokens is generated."""
+  
+  def __init__(self, stops: list = []):
+      super().__init__()
+      self.stops = stops
+  
+  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> bool:
+      return input_ids[0][-len(self.stops) :].tolist() == self.stops
+
+
+stopping_criteria = StoppingCriteriaList(
+  [StoppingCriteriaSub(stops=tokenizer(EOC_FORMAT)["input_ids"])]
+)
+
+
+def predict(input_text: str) -> str:
+  tokenized = tokenizer(
+      input_text,
+      max_length=self.max_length,
+      padding=True,
+      truncation=True,
+      return_tensors="pt",
+  )
+
+  with torch.no_grad():
+      out = model.generate(
+          input_ids=tokenized["input_ids"].to(DEVICE),
+          attention_mask=tokenized["attention_mask"].to(DEVICE),
+          pad_token_id=self.tokenizer.eos_token_id,
+          max_new_tokens=max_new_tokens,
+          num_return_sequences=num_return_sequences,
+          do_sample=True,
+          temperature=temperature,
+          top_p=top_p,
+          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)
+```
+
 ## Sample training data
 
 ```json