-
Notifications
You must be signed in to change notification settings - Fork 239
/
main.py
59 lines (43 loc) · 1.36 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import shutil
import asyncio
import argparse
import ollama
async def speak(speaker, content):
if speaker:
p = await asyncio.create_subprocess_exec(speaker, content)
await p.communicate()
async def main():
parser = argparse.ArgumentParser()
parser.add_argument('--speak', default=False, action='store_true')
args = parser.parse_args()
speaker = None
if not args.speak:
...
elif say := shutil.which('say'):
speaker = say
elif (espeak := shutil.which('espeak')) or (espeak := shutil.which('espeak-ng')):
speaker = espeak
client = ollama.AsyncClient()
messages = []
while True:
if content_in := input('>>> '):
messages.append({'role': 'user', 'content': content_in})
content_out = ''
message = {'role': 'assistant', 'content': ''}
async for response in await client.chat(model='mistral', messages=messages, stream=True):
if response['done']:
messages.append(message)
content = response['message']['content']
print(content, end='', flush=True)
content_out += content
if content in ['.', '!', '?', '\n']:
await speak(speaker, content_out)
content_out = ''
message['content'] += content
if content_out:
await speak(speaker, content_out)
print()
try:
asyncio.run(main())
except (KeyboardInterrupt, EOFError):
...