202 lines
6.8 KiB
Python
202 lines
6.8 KiB
Python
import gi, ollama, asyncio, threading
|
|
gi.require_version("Gtk", "4.0")
|
|
from gi.repository import GLib, Gtk, Gdk
|
|
|
|
import util
|
|
from chat_session import ChatSession
|
|
|
|
@Gtk.Template(filename="pyllama.ui")
|
|
class PyLlamaWindow(Gtk.ApplicationWindow):
|
|
__gtype_name__ = "main_window"
|
|
|
|
model_list = Gtk.Template.Child("model_list")
|
|
model_parameters = Gtk.Template.Child("model_parameters")
|
|
chat_scroll = Gtk.Template.Child("chat_scroll")
|
|
chat_list = Gtk.Template.Child("chat_list")
|
|
prompt = Gtk.Template.Child("prompt")
|
|
prompt_controller = Gtk.Template.Child("prompt_controller")
|
|
send_button = Gtk.Template.Child("send_button")
|
|
cancel_button = Gtk.Template.Child("cancel_button")
|
|
gpu_layers_spin = Gtk.Template.Child("gpu_layers_spin")
|
|
low_vram_switch = Gtk.Template.Child("low_vram_switch")
|
|
threads_spin = Gtk.Template.Child("threads_spin")
|
|
temperature_spin = Gtk.Template.Child("temperature_spin")
|
|
|
|
def __init__(self, application=None, title=None):
|
|
super().__init__(application=application, title=title)
|
|
self.model_list.set_model(self.get_application().model_store)
|
|
self.prompt.add_controller(self.prompt_controller)
|
|
|
|
self._current_response = None
|
|
|
|
@property
|
|
def prompt_buffer(self):
|
|
return self.prompt.get_buffer()
|
|
|
|
def prompt_clear(self):
|
|
self.prompt_buffer.set_text("")
|
|
|
|
def prompt_get_text(self):
|
|
start = self.prompt_buffer.get_start_iter()
|
|
end = self.prompt_buffer.get_end_iter()
|
|
return self.prompt_buffer.get_text(start, end, False).strip()
|
|
|
|
def prompt_hold(self, *args):
|
|
self.grab_focus()
|
|
self.prompt.set_sensitive(False)
|
|
self.send_button.set_visible(False)
|
|
self.cancel_button.set_visible(True)
|
|
|
|
def prompt_unhold(self, *args):
|
|
self.prompt.set_sensitive(True)
|
|
self.prompt.grab_focus()
|
|
self.send_button.set_visible(True)
|
|
self.cancel_button.set_visible(False)
|
|
|
|
def chat_scroll_to_bottom(self):
|
|
adj = self.chat_scroll.get_vadjustment()
|
|
adj.set_value(adj.get_upper() - adj.get_page_size())
|
|
|
|
def chat_start_response(self):
|
|
self._current_response = self.chat_add_message(sender=False)
|
|
|
|
def chat_add_message(self, text="", sender=True):
|
|
label = Gtk.Label(selectable=True, wrap=True, xalign=1.0 if sender else 0.0)
|
|
label.set_text(text)
|
|
util.set_margin_all(label, 8)
|
|
|
|
row = Gtk.ListBoxRow()
|
|
row.set_child(label)
|
|
self.chat_list.append(row)
|
|
|
|
self.chat_scroll_to_bottom()
|
|
return label
|
|
|
|
def chat_add_response_part(self, part):
|
|
self._current_response.set_label(self._current_response.get_label() + part)
|
|
self.chat_scroll_to_bottom()
|
|
|
|
@Gtk.Template.Callback()
|
|
def prompt_key_pressed(self, _a, keyval, _b, state):
|
|
match keyval:
|
|
case Gdk.KEY_Return if not state & Gdk.ModifierType.SHIFT_MASK:
|
|
self.send_button_clicked(self.send_button)
|
|
return True
|
|
case _:
|
|
return False
|
|
|
|
@Gtk.Template.Callback()
|
|
def send_button_clicked(self, button):
|
|
text = self.prompt_get_text()
|
|
self.prompt_clear()
|
|
|
|
self.chat_add_message(text, sender=True)
|
|
self.chat_start_response()
|
|
|
|
app = self.get_application()
|
|
model = self.model_list.get_selected_item().get_string()
|
|
|
|
app.history_push(role="user", text=text)
|
|
|
|
options = {
|
|
"num_gpu": self.gpu_layers_spin.get_value_as_int(),
|
|
"low_vram": self.low_vram_switch.get_state(),
|
|
"num_thread": self.threads_spin.get_value_as_int(),
|
|
"temperature": self.temperature_spin.get_value_as_int()
|
|
}
|
|
|
|
async def task():
|
|
await app.send_prompt(model, options)
|
|
|
|
if app.current_async_task and not app.current_async_task.done():
|
|
app.current_async_task.cancel()
|
|
|
|
app.current_async_task = asyncio.run_coroutine_threadsafe(task(), app.async_loop)
|
|
|
|
@Gtk.Template.Callback()
|
|
def cancel_button_clicked(self, *args):
|
|
app = self.get_application()
|
|
app.stop_generation()
|
|
|
|
@Gtk.Template.Callback()
|
|
def gpu_layers_spin_changed(self, *args):
|
|
pass
|
|
|
|
@Gtk.Template.Callback()
|
|
def low_vram_switch_set(self, *args):
|
|
pass
|
|
|
|
@Gtk.Template.Callback()
|
|
def threads_spin_changed(self, *args):
|
|
pass
|
|
|
|
@Gtk.Template.Callback()
|
|
def temperature_spin_changed(self, *args):
|
|
pass
|
|
|
|
class PyLlama(Gtk.Application):
|
|
def __init__(self):
|
|
super().__init__(application_id="eur.cfpi-fpsi.pyllama")
|
|
GLib.set_application_name("PyLlama")
|
|
|
|
self.num_gpu = 15
|
|
self.history_size = 4
|
|
self.history = []
|
|
self.model_store = Gtk.StringList()
|
|
self._populate_models()
|
|
|
|
self.client = ollama.AsyncClient()
|
|
self.async_loop = asyncio.new_event_loop()
|
|
self.current_async_task = None
|
|
threading.Thread(target=self._run_asyncio_loop, daemon=True).start()
|
|
|
|
def do_activate(self):
|
|
self.window = PyLlamaWindow(application=self, title="PyLlama")
|
|
self.window.present()
|
|
|
|
def _run_asyncio_loop(self):
|
|
asyncio.set_event_loop(self.async_loop)
|
|
self.async_loop.run_forever()
|
|
return True
|
|
|
|
def _populate_models(self):
|
|
try:
|
|
models = ollama.list()
|
|
names = [m.model for m in models["models"]]
|
|
except Exception as e:
|
|
print ("Failed to fetch models: ", e)
|
|
names = []
|
|
|
|
self.model_store.splice(0, len(self.model_store), names)
|
|
|
|
def stop_generation(self):
|
|
if self.current_async_task and not self.current_async_task.done():
|
|
self.current_async_task.cancel()
|
|
|
|
def history_push(self, role, text):
|
|
self.history.append({"role": role, "content": text})
|
|
while len(self.history) > self.history_size:
|
|
self.history.pop(0)
|
|
|
|
async def send_prompt(self, model, options):
|
|
GLib.idle_add(self.window.prompt_hold, True)
|
|
try:
|
|
messages = list(self.history)[-self.history_size:]
|
|
async for part in await self.client.chat(model=model, messages=messages, stream=True, options=options):
|
|
chunk = part.message.content
|
|
|
|
if not self.history or self.history[-1]["role"] != "assistant":
|
|
self.history.append({"role": "assistant", "content": chunk})
|
|
else:
|
|
self.history[-1]["content"] += chunk
|
|
|
|
GLib.idle_add(self.window.chat_add_response_part, part.message.content)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
finally:
|
|
GLib.idle_add(self.window.prompt_unhold, True)
|
|
|
|
if __name__ == "__main__":
|
|
app = PyLlama()
|
|
app.run()
|