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}
}
```