File size: 2,326 Bytes
684bce1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
---
license: cc
language:
- en
---
This is the proposition segmentation model from "Sub-Sentence Encoder: Contrastive Learning of Propositional Semantic Representations" by Chen et. al. 2023.
## Usage
The prompt to the model is formatted like: `segment sentence: {input_sentence}`.
For each sentence, the model will output the propositions concatenated by `[sep]` as a string.
For example, if we use the following example code to segment `"Dracula is a novel by Bram Stoker featuring Count Dracula as the protagonist."`.
The model output will be `['Dracula is a novel by Bram Stoker.[sep]Count Dracula is the protagonist of Dracula.']`
```
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
gen_kwargs = {
"length_penalty": 0,
"max_new_tokens": 256,
"min_length": 10,
"no_repeat_ngram_size": 0,
"num_beams": 1,
}
SEGMENT5_PROMPT = "segment sentence: {}"
SEGMENT5_SEP_TOKEN = "[sep]"
model = AutoModelForSeq2SeqLM.from_pretrained("sihaochen/SegmenT5-large")
tokenizer = AutoTokenizer.from_pretrained("sihaochen/SegmenT5-large")
model.eval()
# define an example input sentence
example_sentence = "Dracula is a novel by Bram Stoker featuring Count Dracula as the protagonist."
example_input = SEGMENT5_PROMPT.format(example_sentence)
input_ids = tokenizer(example_input,
return_tensors="pt",
padding="max_length",
max_length=512,
truncation=True).input_ids
logits = model.generate(input_ids, **gen_kwargs)
outputs = tokenizer.batch_decode(logits, skip_special_tokens=True)
output = outputs[0].split(SEGMENT5_SEP_TOKEN)
print(output)
# Output: ['Dracula is a novel by Bram Stoker.', 'Count Dracula is the protagonist of Dracula.']
```
## Sub-Sentence Encoder
For model checkpoints + code for the sub-sentence encoders, checkout: https://github.com/schen149/sub-sentence-encoder/
## Citation
```
@article{chen2023subsentence,
title={Sub-Sentence Encoder: Contrastive Learning of Propositional Semantic Representations},
author={Sihao Chen and Hongming Zhang and Tong Chen and Ben Zhou and Wenhao Yu and Dian Yu and Baolin Peng and Hongwei Wang and Dan Roth and Dong Yu},
journal={arXiv preprint arXiv:2311.04335},
year={2023},
URL = {https://arxiv.org/pdf/2311.04335.pdf}
}
```
|