Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DNM: Patch FT Tagger #210

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

DNM: Patch FT Tagger #210

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1a6aec7
DNM: Patch FT Tagger
undfined Sep 25, 2024
ba87b1c
Tag og text
undfined Oct 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions python/dolma/core/ft_tagger.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@
from fasttext import train_supervised
from fasttext.FastText import _FastText

from .data_types import DocResult, Document, Span, TextSlice
from .taggers import BaseTagger
from .data_types import DocResult, Document, DocumentWithMetadata, Span, TextSlice
from .taggers import BaseTaggerWithMetadata
from .utils import split_paragraphs, split_sentences


Expand All @@ -25,7 +25,7 @@ class Prediction(NamedTuple):
score: float


class BaseFastTextTagger(BaseTagger):
class BaseFastTextTagger(BaseTaggerWithMetadata):
SENTENCE_LEVEL_TAGGER = "sentence"
PARAGRAPH_LEVEL_TAGGER = "paragraph"
DOCUMENT_LEVEL_TAGGER = "document"
Expand Down Expand Up @@ -135,13 +135,13 @@ def test(
model_performance = classifier.test(local_test_file)
print(model_performance)

def predict(self, doc: Document) -> DocResult:
def predict(self, doc: DocumentWithMetadata) -> DocResult:
if self.mode == self.SENTENCE_LEVEL_TAGGER:
units = split_sentences(doc.text)
elif self.mode == self.PARAGRAPH_LEVEL_TAGGER:
units = split_paragraphs(doc.text)
elif self.mode == self.DOCUMENT_LEVEL_TAGGER:
units = [TextSlice(doc=doc.text, start=0, end=len(doc.text))]
units = [TextSlice(doc=doc.metadata["original_text"], start=0, end=len(doc.metadata["original_text"]))]
else:
raise ValueError(f"Unknown mode {self.mode}")

Expand Down
13 changes: 7 additions & 6 deletions python/dolma/taggers/length.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@
import uniseg.wordbreak
from tokenizers import Regex, Tokenizer, pre_tokenizers

from ..core.data_types import DocResult, Document, Span, TextSlice
from ..core.data_types import DocResult, Document, DocumentWithMetadata, Span, TextSlice
from ..core.registry import TaggerRegistry
from ..core.taggers import BaseTagger
from ..core.taggers import BaseTagger, BaseTaggerWithMetadata
from ..core.utils import split_paragraphs


Expand Down Expand Up @@ -161,16 +161,17 @@ def predict(self, doc: Document) -> DocResult:


@TaggerRegistry.add("dolma_v1_tokenizer")
class DolmaV1Tokenizer(BaseTagger):
class DolmaV1Tokenizer(BaseTaggerWithMetadata):
TOKENIZER_NAME_OR_PATH = "allenai/gpt-neox-olmo-dolma-v1_5"

def __init__(self) -> None:
self.tokenizer = Tokenizer.from_pretrained(self.TOKENIZER_NAME_OR_PATH)
super().__init__()

def predict(self, doc: Document) -> DocResult:
score = len(self.tokenizer.encode(text)) if (text := doc.text.strip()) else 0
return DocResult(doc=doc, spans=[Span(start=0, end=len(doc.text), type="length", score=score)])
def predict(self, doc: DocumentWithMetadata) -> DocResult:
text = doc.metadata["original_text"].strip()
score = len(self.tokenizer.encode(text)) if (text) else 0
return DocResult(doc=doc, spans=[Span(start=0, end=len(text), type="length", score=score)])


@TaggerRegistry.add("dolma_v2_tokenizer")
Expand Down
Loading