from transformers import pipeline
import torch

gec_pipe = pipeline(
    "text-generation",
    model="tartuNLP/Llammas-base-p1-llama-errors-p2-GEC",
    torch_dtype=torch.bfloat16,
    device_map="auto",
    do_sample=False, num_beams=4, temperature=None, top_p=None
)
gec_pipe.tokenizer.pad_token_id = gec_pipe.tokenizer.eos_token_id
gec_pipe.tokenizer.padding_side = "left"

### Input sentence here:
input_sentence = "Ma läheb koju"
gec_pipe([{"role": "user", "content": input_sentence}], max_new_tokens=300)[0]["generated_text"][-1]["content"]

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained(
    "tartuNLP/Llammas-base-p1-llama-errors-p2-GEC",
    device_map="auto",
    return_dict=True,
    low_cpu_mem_usage=True,
    torch_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(
    "tartuNLP/Llammas-base-p1-llama-errors-p2-GEC",
    padding_side="left"
)
# Need to set the padding token to 0 or eos_token_id if batching is used
# (the model does not set it by default)
tokenizer.pad_token_id = tokenizer.eos_token_id

gec_pipe = pipeline(
    "text-generation", model=model, tokenizer=tokenizer, do_sample=False, num_beams=4, temperature=None, top_p=None
)


### Input sentence here
input_sentence = "Ma läheb koju"

# Two options:
# 1)
PROMPT = '### Instruction:\nReply with a corrected version of the input sentence in Estonian with all grammatical and spelling errors fixed. If there are no errors, reply with a copy of the original sentence.\n\n### Input:\n{input}\n\n### Response:\n'
example = PROMPT.format(input=input_sentence)
# 2) or use the chat template provided by us that does the same thing
example = tokenizer.apply_chat_template([{"role": "user", "content": input_sentence}], tokenize=False)

gec_pipe(example, max_new_tokens=300)[0]["generated_text"][len(example):]

@inproceedings{luhtaru-etal-2024-err,
    title = "To Err Is Human, but Llamas Can Learn It Too",
    author = "Luhtaru, Agnes  and
      Purason, Taido  and
      Vainikko, Martin  and
      Del, Maksym  and
      Fishel, Mark",
    editor = "Al-Onaizan, Yaser  and
      Bansal, Mohit  and
      Chen, Yun-Nung",
    booktitle = "Findings of the Association for Computational Linguistics: EMNLP 2024",
    month = nov,
    year = "2024",
    address = "Miami, Florida, USA",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.findings-emnlp.727",
    doi = "10.18653/v1/2024.findings-emnlp.727",
    pages = "12466--12481",
    abstract = "This study explores enhancing grammatical error correction (GEC) through automatic error generation (AEG) using language models (LMs). Specifically, we fine-tune Llama 2 LMs for error generation and find that this approach yields synthetic errors akin to human errors. Next, we train GEC Llama models using these artificial errors and outperform previous state-of-the-art error correction models, with gains ranging between 0.8 and 6 F0.5 points across all tested languages (German, Ukrainian, and Estonian). Moreover, we demonstrate that generating errors by fine-tuning smaller sequence-to-sequence models and prompting large commercial LMs (GPT3.5 and GPT4) also results in synthetic errors beneficially affecting error generation models. We openly release trained models for error generation and correction as well as all the synthesized error datasets for the covered languages.",
}