Update README.md

README.md

@@ -17,6 +17,22 @@ metrics:
 > reading, video games, walking
 
 ### How to use in code
-`from transformers import AutoTokenizer, AutoModelForSeq2SeqLM`\
-`tokenizer = AutoTokenizer.from_pretrained("njvdnbus/Interest_extraction")`\
-`model = AutoModelForSeq2SeqLM.from_pretrained("njvdnbus/Interest_extraction")`
+```{python}
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+tokenizer = AutoTokenizer.from_pretrained("njvdnbus/Interest_extraction")
+model = AutoModelForSeq2SeqLM.from_pretrained("njvdnbus/Interest_extraction")
+
+def use_model(text):
+    inputs = ["summarize: " + text]
+    inputs = tokenizer(inputs, truncation=True, return_tensors="pt")
+    output = model.generate(**inputs, num_beams=1, do_sample=True, min_length=1, max_length=64)
+    decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    predicted_interests = nltk.sent_tokenize(decoded_output.strip())[0]
+    return predicted_interests
+
+text= '''What other hobbies do you have?
+When I have time I like to cook for my family. Most often this only happens in the weekends.'''
+print(use_model(text))
+```
+
+> cooking