|
10 | 10 |
|
11 | 11 | from licensedcode.tokenize import required_phrase_splitter |
12 | 12 |
|
| 13 | +from scancode_required_phrases.training import encode_complete_words |
13 | 14 | from scancode_required_phrases.training import extract_spans |
14 | 15 | from scancode_required_phrases.training import first_subword_positions |
15 | 16 | from scancode_required_phrases.training import ID2LABEL |
@@ -41,41 +42,6 @@ def words_from_text(text): |
41 | 42 | return required_phrase_splitter(unicodedata.normalize("NFKC", text)) |
42 | 43 |
|
43 | 44 |
|
44 | | -def _word_counts(word_ids): |
45 | | - counts = {} |
46 | | - for word_id in word_ids: |
47 | | - if word_id is not None: |
48 | | - counts[word_id] = counts.get(word_id, 0) + 1 |
49 | | - return counts |
50 | | - |
51 | | - |
52 | | -def encode_words(tokenizer, words, max_length): |
53 | | - """Encode the longest complete-word prefix and report truncation.""" |
54 | | - call = dict( |
55 | | - is_split_into_words=True, |
56 | | - add_special_tokens=True, |
57 | | - return_tensors="pt", |
58 | | - ) |
59 | | - full = tokenizer(words, truncation=False, **call) |
60 | | - encoding = tokenizer(words, truncation=True, max_length=max_length, **call) |
61 | | - |
62 | | - full_counts = _word_counts(full.word_ids()) |
63 | | - retained_counts = _word_counts(encoding.word_ids()) |
64 | | - covered_words = max(retained_counts, default=-1) + 1 |
65 | | - complete_words = covered_words |
66 | | - |
67 | | - if covered_words and retained_counts[covered_words - 1] != full_counts[covered_words - 1]: |
68 | | - complete_words -= 1 |
69 | | - encoding = tokenizer(words[:complete_words], truncation=False, **call) |
70 | | - |
71 | | - if not complete_words: |
72 | | - raise ValueError("Tokenizer retained no complete words") |
73 | | - if encoding["input_ids"].shape[1] > max_length: |
74 | | - raise ValueError("Complete-word encoding exceeds the model maximum length") |
75 | | - |
76 | | - return encoding, complete_words < len(words) |
77 | | - |
78 | | - |
79 | 45 | def span_confidence(crf, word_emissions, tags, mask, free, span): |
80 | 46 | """Return the CRF probability mass agreeing with one decoded span.""" |
81 | 47 | start, end = span |
@@ -117,18 +83,22 @@ def predict(self, text): |
117 | 83 | if not words: |
118 | 84 | return PredictionResult(words=(), phrases=(), truncated=False) |
119 | 85 |
|
120 | | - encoding, truncated = encode_words( |
| 86 | + encoding, truncated = encode_complete_words( |
| 87 | + tokens=words, |
121 | 88 | tokenizer=self.tokenizer, |
122 | | - words=words, |
123 | 89 | max_length=self.max_length, |
124 | 90 | ) |
125 | 91 | positions = first_subword_positions(encoding.word_ids()) |
126 | 92 | if not positions: |
127 | 93 | return PredictionResult(words=tuple(words), phrases=(), truncated=False) |
128 | 94 |
|
129 | 95 | device = next(self.model.parameters()).device |
130 | | - input_ids = encoding["input_ids"].to(device) |
131 | | - attention_mask = encoding["attention_mask"].to(device) |
| 96 | + input_ids = torch.tensor([encoding["input_ids"]], dtype=torch.long, device=device) |
| 97 | + attention_mask = torch.tensor( |
| 98 | + [encoding["attention_mask"]], |
| 99 | + dtype=torch.long, |
| 100 | + device=device, |
| 101 | + ) |
132 | 102 |
|
133 | 103 | with torch.inference_mode(): |
134 | 104 | emissions = self.model.emissions(input_ids, attention_mask) |
|
0 commit comments