├── LICENSE ├── README.md └── main.py /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 Patrick Loeber 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AI powered typing assistant with Ollama 2 | 3 | A script that can run in the background and listen to hotkeys, then uses a Large Language Model to fix the text. Less than 100 lines of code. 4 | 5 | Inspired by this tweet: 6 | 7 | https://twitter.com/karpathy/status/1725553780878647482 8 | 9 | > "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 10 | 11 | You'll find a demo and step-by-step code explanations on my YouTube channel: 12 | 13 | [![Alt text](https://img.youtube.com/vi/IUTFrexghsQ/hqdefault.jpg)](https://youtu.be/IUTFrexghsQ) 14 | 15 | ## Get Started 16 | 17 | ### 1. Set up Ollama 18 | 19 | Ollama Installation: https://github.com/ollama/ollama 20 | 21 | Run `ollama run mistral:7b-instruct-v0.2-q4_K_S` 22 | 23 | Mistal 7B Instruct works well for this task, but feel free to try other models, too :) 24 | 25 | ### 2. Install dependencies 26 | ``` 27 | pip install pynput pyperclip httpx 28 | ``` 29 | 30 | - pynput: https://pynput.readthedocs.io/en/latest/ 31 | - pyperclip: https://github.com/asweigart/pyperclip 32 | - httpx: https://github.com/encode/httpx/ 33 | 34 | ### 3. Run it 35 | 36 | Start the assistant: 37 | 38 | ``` 39 | python main.py 40 | ``` 41 | 42 | Hotkeys you can then press: 43 | 44 | - F9: Fixes the current line (without having to select the text) 45 | - F10: Fixes the current selection 46 | 47 | **Note**: You may get an error the first time saying `This process is not trusted! Input event monitoring will not be possible until it is added to accessibility clients`. On Mac, you need to add the script (IDE/terminal) both on accessibility and input monitoring. 48 | 49 | **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`). 50 | 51 | ## Customize 52 | 53 | Hotkeys, prompt, and Ollama config can be easily customized and extended in the code. 54 | 55 | For example, here are some prompt templates you can try: 56 | 57 | ```python 58 | from string import Template 59 | 60 | PROMPT_TEMPLATE_FIX_TEXT = Template( 61 | """Fix all typos and casing and punctuation in this text, but preserve all new line characters: 62 | 63 | $text 64 | 65 | Return only the corrected text.""" 66 | ) 67 | 68 | PROMPT_TEMPLATE_GENERATE_TEXT = Template( 69 | """Generate a snarky paragraph with 3 sentences about the following topic: 70 | 71 | $text 72 | 73 | Return only the corrected text.""" 74 | ) 75 | 76 | PROMPT_TEMPLATE_SUMMARIZE_TEXT = Template( 77 | """Summarize the following text in 3 sentences: 78 | 79 | $text 80 | 81 | Return only the corrected text.""" 82 | ) 83 | ``` 84 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import time 2 | from string import Template 3 | 4 | import httpx 5 | from pynput import keyboard 6 | from pynput.keyboard import Key, Controller 7 | import pyperclip 8 | 9 | 10 | controller = Controller() 11 | 12 | OLLAMA_ENDPOINT = "http://localhost:11434/api/generate" 13 | OLLAMA_CONFIG = { 14 | "model": "mistral:7b-instruct-v0.2-q4_K_S", 15 | "keep_alive": "5m", 16 | "stream": False, 17 | } 18 | 19 | PROMPT_TEMPLATE = Template( 20 | """Fix all typos and casing and punctuation in this text, but preserve all new line characters: 21 | 22 | $text 23 | 24 | Return only the corrected text, don't include a preamble. 25 | """ 26 | ) 27 | 28 | 29 | def fix_text(text): 30 | prompt = PROMPT_TEMPLATE.substitute(text=text) 31 | response = httpx.post( 32 | OLLAMA_ENDPOINT, 33 | json={"prompt": prompt, **OLLAMA_CONFIG}, 34 | headers={"Content-Type": "application/json"}, 35 | timeout=10, 36 | ) 37 | if response.status_code != 200: 38 | print("Error", response.status_code) 39 | return None 40 | return response.json()["response"].strip() 41 | 42 | 43 | def fix_current_line(): 44 | # macOS short cut to select current line: Cmd+Shift+Left 45 | controller.press(Key.cmd) 46 | controller.press(Key.shift) 47 | controller.press(Key.left) 48 | 49 | controller.release(Key.cmd) 50 | controller.release(Key.shift) 51 | controller.release(Key.left) 52 | 53 | fix_selection() 54 | 55 | 56 | def fix_selection(): 57 | # 1. Copy selection to clipboard 58 | with controller.pressed(Key.cmd): 59 | controller.tap("c") 60 | 61 | # 2. Get the clipboard string 62 | time.sleep(0.1) 63 | text = pyperclip.paste() 64 | 65 | # 3. Fix string 66 | if not text: 67 | return 68 | fixed_text = fix_text(text) 69 | if not fixed_text: 70 | return 71 | 72 | # 4. Paste the fixed string to the clipboard 73 | pyperclip.copy(fixed_text) 74 | time.sleep(0.1) 75 | 76 | # 5. Paste the clipboard and replace the selected text 77 | with controller.pressed(Key.cmd): 78 | controller.tap("v") 79 | 80 | 81 | def on_f9(): 82 | fix_current_line() 83 | 84 | 85 | def on_f10(): 86 | fix_selection() 87 | 88 | 89 | with keyboard.GlobalHotKeys({"<101>": on_f9, "<109>": on_f10}) as h: 90 | h.join() 91 | --------------------------------------------------------------------------------