Interest_extraction / README.md
njvdnbus's picture
Update README.md
8517e80
|
raw
history blame
No virus
1.2 kB
metadata
language:
  - en
metrics:
  - f1

Interest extraction

Extracts the interests from a question-answer pair.

Model input

summarize: [QUESTION]
[ANSWER]

Example

summarize: What do you like to do in the weekend?
I like to spend my free time reading, playing video games, and going on walks.

Output

reading, video games, walking

How to use in code

import nltk
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