Build a typing assistant with Ollama, Mistral 7B, and Python

Learn how to build an AI-powered locally running typing assistant.

Hi there👋

It’s been a while, but now I’m back with a cool project and tutorial: Learn how to build an AI-powered locally running typing assistant with Ollama and Mistral 7B.

In this post, I'll show you how to do it.

It's a script with less than 100 lines of code that can run in the background and listen to hotkeys, then uses a Large Language Model to fix the text.

It’s inspired by this tweet:

"GPT is surprisingly good at correcting minor typos, so you can write really really fast, ignore mistakes and keep going, and it comes out just fine." - Andrej Karpathy

The full code is on GitHub: https://github.com/patrickloeber/ai-typing-assistant

And here's a demo and a step-by-step coding tutorial if you prefer video content:

Project setup:

Ollama Installation: https://github.com/ollama/ollama. Then run

ollama run mistral:7b-instruct-v0.2-q4_K_S.

I found that Mistal 7B Instruct (quantized) works well for this task, but feel free to try other models, too :)

It relies on those dependencies:

- pynput: To listen to hotkeys and control the keyboard

- pyperclip: Clipboard functions to copy and paste text

- httpx: To send a POST request to Ollama running on your localhost

Install the dependencies with pip:

pip install pynput pyperclip httpx

Everything should work cross-platform, only the shortcuts need to be adapted for your system.

Hotkey listener in Python

This is how you can write a Hotkey listener with Python and pynput:

from pynput import keyboard

def on_f9():

    fix_current_line()


def on_f10():

    fix_selection()


with keyboard.GlobalHotKeys({"<101>": on_f9, "<109>": on_f10}) as h:

    h.join()

This runs endlessly in a separate thread until you stop the script.

Modify lines and selected text with Python

To modify selected lines, we only need the functionality to copy and paste text (= Cmd+C and Cmd+V), and to access and modify the clipboard with Python. For this, set up a keyboard controller with pynput, and use pyperclip for the clipboard functions:

import time

from pynput.keyboard import Key, Controller

import pyperclip


controller = Controller()


def fix_current_line():

    # macOS short cut to select current line: Cmd+Shift+Left

    controller.press(Key.cmd)

    controller.press(Key.shift)

    controller.press(Key.left)

    controller.release(Key.cmd)

    controller.release(Key.shift)

    controller.release(Key.left)


    fix_selection()


def fix_selection():

    # 1. Copy selection to clipboard

    with controller.pressed(Key.cmd):

        controller.tap("c")


    # 2. Get the clipboard string

    time.sleep(0.1)

    text = pyperclip.paste()


    # 3. Fix string

    if not text:

        return

    fixed_text = fix_text(text)

    if not fixed_text:

        return


    # 4. Paste the fixed string to the clipboard

    pyperclip.copy(fixed_text)

    time.sleep(0.1)


    # 5. Paste the clipboard and replace the selected text

    with controller.pressed(Key.cmd):

        controller.tap("v")

Use Ollama and Mistral 7B to fix text

Ollama comes with a REST API that's running on your localhost out of the box. So after the installation and downloading of the model, we only need to implement logic to send a POST request.

We also use a simple prompt that tells the model to fix all typos, casing, and punctuation:

import httpx

OLLAMA_ENDPOINT = "http://localhost:11434/api/generate"

OLLAMA_CONFIG = {

    "model": "mistral:7b-instruct-v0.2-q4_K_S",

    "keep_alive": "5m",

    "stream": False,

}


PROMPT_TEMPLATE = Template(

    """Fix all typos and casing and punctuation in this text, but preserve all new line characters:

$text

Return only the corrected text, don't include a preamble.

"""

)


def fix_text(text):

    prompt = PROMPT_TEMPLATE.substitute(text=text)

    response = httpx.post(

        OLLAMA_ENDPOINT,

        json={"prompt": prompt, **OLLAMA_CONFIG},

        headers={"Content-Type": "application/json"},

        timeout=10,

    )

    if response.status_code != 200:

        print("Error", response.status_code)

        return None

    return response.json()["response"].strip()

Run the typing assistant

Put everything together and start the assistant:

python main.py

Hotkeys you can then press:

- F9: Fixes the current line (without having to select the text)

- F10: Fixes the current selection

Note: The code works on macOS. The underlying shortcuts in the code like Cmd+Shift+Left, Cmd+C, Cmd+V might have to be changed for Linux & Windows (e.g. Key.cmd -> Key.ctrl).

Customize

The hotkeys, prompt, and Ollama config can be easily customized and extended in the code.

You can easily add hotkeys that summarize or generate new text.

For example, here are some fun other prompt templates you can try:

from string import Template


PROMPT_TEMPLATE_FIX_TEXT = Template(

    """Fix all typos and casing and punctuation in this text, but preserve all new line characters:

$text

Return only the corrected text."""

)


PROMPT_TEMPLATE_GENERATE_TEXT = Template(

    """Generate a snarky paragraph with 3 sentences about the following topic:

$text

Return only the corrected text."""

)


PROMPT_TEMPLATE_SUMMARIZE_TEXT = Template(

    """Summarize the following text in 3 sentences:

$text

Return only the corrected text."""

)

Have fun with it!