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

I built a locally running typing assistant with Ollama, Mistral 7B, and Python. 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.

Inspired by this tweet:

https://twitter.com/karpathy/status/1725553780878647482

"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:

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

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.

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.

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")

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()

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).

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!