examples/use_with_transformers.py

"""Use EDEN through the Hugging Face Transformers integration.

Run from the repository root:

    python examples/use_with_transformers.py
"""

import os
import sys

# Make the repository root importable so the custom model code resolves when
# loading from this local folder. When loading from the Hugging Face Hub this is
# handled for you and you can simply use the model id, for example "Rybib/EDEN".
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO_ROOT)

from transformers import AutoModel, AutoTokenizer

# Point this at the published model id or a local path to this repository.
MODEL_ID = REPO_ROOT

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).eval()

examples = [
    "i relly wnt this sentance to sound more profesional",
    "me and him goes to the store yesterday and buyed some thing",
    "their are alot of reasons why this dont work proper",
]

for rough in examples:
    polished = model.enhance(tokenizer, rough, strategy="beam")
    print("rough   :", rough)
    print("polished:", polished)
    print()