Text proofreading guide for Python

The MediaPipe Text Proofreader task lets you identify spelling, grammatical, and stylistic errors in a text and generate a corrected version while maintaining the original meaning. These instructions show you how to use the Text Proofreader with Python.

For more information about the capabilities, models, and configuration options of this task, see the Overview.

Code example

The example code for Text Proofreader provides a complete implementation of this task in Python for your reference. This code helps you test this task and get started on building your own text proofreader. You can view the source code for this example on GitHub.

Setup

This section describes key steps for setting up your development environment and code projects specifically to use Text Proofreader. For general information on setting up your development environment for using MediaPipe tasks, including platform version requirements, see the Setup guide for Python.

Packages

Text Proofreader uses the mediapipe pip package. You can install the dependency with the following command:

$ python -m pip install mediapipe

Imports

Import the following classes to access the Text Proofreader task functions:

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import text

Model

The MediaPipe Text Proofreader task requires a trained model that is compatible with this task. For more information on available trained models for Text Proofreader, see the task overview Models section.

Select and download a model, and then store it in a local directory.

model_path = '/absolute/path/to/proofreader.litertlm'

Specify the path of the model within the model_asset_path parameter, as shown below:

base_options = python.BaseOptions(model_asset_path=model_path)

Create the task

The MediaPipe Text Proofreader task uses the create_from_options function to set up the task. The create_from_options function accepts values for configuration options. You can also initialize the task using the create_from_model_path factory function. The create_from_model_path function accepts a relative or absolute path to the trained model file.

For more information on configuration options, see Configuration options.

The following code demonstrates how to build and configure this task.

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
TextProofreader = mp.tasks.text.TextProofreader
TextProofreaderOptions = mp.tasks.text.TextProofreaderOptions
TextProofreaderResult = mp.tasks.text.TextProofreaderResult
TextProofreaderCorrection = mp.tasks.text.TextProofreaderCorrection
TextProofreaderCorrectionType = mp.tasks.text.TextProofreaderCorrectionType

# Build the configuration options
options = TextProofreaderOptions(
    base_options=BaseOptions(model_asset_path=model_path),
    max_num_tokens=4096
)

text_proofreader = TextProofreader.create_from_options(options)

Configuration options

This task has the following configuration options for Python applications:

Option Name Description Value Range / Type Default Value
max_num_tokens The maximum number of tokens for proofread tasks. If set, the proofread output will be truncated if the input and output exceed this value. If not set (or equal to 0), the default limit is decided by the model capacity (8k). Integer None

Prepare data

Text Proofreader works with text (str) data. The task handles the data input preprocessing, including tokenization and tensor preprocessing. You do not need to perform any additional preprocessing of the input text beforehand.

input_text = "The long input text to be proofread goes here..."

Run the task

The Text Proofreader uses the proofread function for synchronous execution, and proofread_async for asynchronous streaming execution to generate corrected text.

Synchronous Execution

The following code demonstrates how to execute the processing synchronously.

# Perform text proofreading on the provided input text.
proofreading_result = text_proofreader.proofread(input_text)

Asynchronous Streaming Execution

For executing asynchronous streaming where parts of the text are generated chunk by chunk, use proofread_async. You must provide a callback function to handle these incoming results.

from typing import Optional

def result_callback(
    result: Optional[mp.tasks.text.TextProofreaderResult],
    error: Optional[str]
):
    if error:
        print(f"Error: {error}")
        return

    if result and result.proofread_text:
        # Result contains the partial chunk string
        print(result.proofread_text, end="")

    if result and result.done:
        print("\nFinished stream!")
        # Access corrections here if needed:
        # print(result.corrections)

# Perform streaming text proofreading on the provided input text.
text_proofreader.proofread_async(input_text, result_callback)

Handle and display results

The Text Proofreader outputs a TextProofreaderResult that contains the completed or partial proofed text, a list of TextProofreaderCorrection objects, and a boolean flag indicating completion status. Each TextProofreaderCorrection specifies a TextProofreaderCorrectionType (such as SAME, INSERTION, or DELETION) and the associated text string.

For streaming tasks, the python TextProofreaderResult maps the package updates with the newly appended token chunk populated into the proofread_text property, and the done property indicating completion.

The following shows an example of the output data from this task:

# Print the output proofread text
print("Proofread Text:", proofreading_result.proofread_text)

# Iterate over granular corrections to understand the diff
for correction in proofreading_result.corrections:
    print(f"[{correction.type.name}] {correction.text}")