diff --git a/main.py b/docman.py similarity index 65% rename from main.py rename to docman.py index 8a3df53..833cbe9 100644 --- a/main.py +++ b/docman.py @@ -13,23 +13,30 @@ class PyLlamaWindow(Gtk.ApplicationWindow): super().__init__(application=application, title=title) self.set_default_size(800, 600) self.cur_ai_response = None + self.gpu_spin = None self.model_store = Gtk.StringList() self.model_dropdown = Gtk.DropDown() self.model_dropdown.set_model(self.model_store) self.model_dropdown.set_tooltip_text("Select model") - self.model_dropdown.set_enable_search(True) + + model_preferences = Gtk.Button() + icon = Gtk.Image.new_from_icon_name("view-more-symbolic") + model_preferences.set_child(icon) + model_preferences.set_tooltip_text("Preferences") + model_preferences.connect("clicked", self.on_preferences_clicked) header = Gtk.HeaderBar() header.set_show_title_buttons(True) header.pack_start(self.model_dropdown) + header.pack_start(model_preferences) self.set_titlebar(header) - vbox = Gtk.Box.new(Gtk.Orientation.VERTICAL, 6) + vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6) set_margin_all(vbox, 8) - self.listbox = Gtk.ListBox.new() + self.listbox = Gtk.ListBox() self.listbox.set_selection_mode(Gtk.SelectionMode.NONE) self.listbox.set_focusable(True) @@ -54,6 +61,35 @@ class PyLlamaWindow(Gtk.ApplicationWindow): self.current_model = None self.populate_models() + def hold_prompt(self, *_): + self.entry.set_placeholder_text("Wait for response...") + self.listbox.grab_focus() + self.entry.set_sensitive(False) + + def unhold_prompt(self, *_): + self.entry.set_placeholder_text("Type message and press Enter...") + self.entry.set_sensitive(True) + self.entry.grab_focus() + + def on_preferences_clicked(self, button): + box = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6) + set_margin_all(box, 8) + + label = Gtk.Label(label="GPU Layers") + label.set_xalign(0.0) + + self.gpu_spin = Gtk.SpinButton(adjustment=Gtk.Adjustment( + value=10, lower=0, upper=100, step_increment=1 + ), digits=0, hexpand=True) + + box.append(label) + box.append(self.gpu_spin) + + preferences_popover = Gtk.Popover(position=Gtk.PositionType.BOTTOM) + preferences_popover.set_child(box) + preferences_popover.set_parent(button) + preferences_popover.popup() + def populate_models(self): try: models = ollama.list() @@ -64,18 +100,19 @@ class PyLlamaWindow(Gtk.ApplicationWindow): names = [m.model for m in models["models"]] self.model_store.splice(0, len(list(self.model_store)), names) - def add_message(self, text=None, sender=True): + def add_message(self, text="", sender=True): label = Gtk.Label() label.set_selectable(True) label.set_wrap(True) label.set_xalign(1.0 if sender else 0.0) - label.set_text("" if not text else text) + label.set_text(text) set_margin_all(label, 8) - + row = Gtk.ListBoxRow() row.set_child(label) - self.listbox.append(row) + self.listbox.append(row) + adj = self.listbox.get_adjustment() adj.set_value(adj.get_upper()) @@ -95,15 +132,16 @@ class PyLlamaWindow(Gtk.ApplicationWindow): buffer.delete_text(0, -1) # Delete all text self.add_message(text=text) - self.cur_ai_response = self.add_message(text=None, sender=False) + self.cur_ai_response = self.add_message(sender=False) self.listbox.grab_focus() self.entry.set_sensitive(False) app = self.get_application() model = self.model_dropdown.get_selected_item().get_string() + num_gpu = self.gpu_spin.get_value_as_int() if self.gpu_spin else 10 app.loop.call_soon_threadsafe(lambda: asyncio.create_task( - app.send_prompt(model, text) + app.send_prompt(model, text, num_gpu) )) class PyLlama(Gtk.Application): @@ -124,11 +162,12 @@ class PyLlama(Gtk.Application): self.loop.run_forever() return True - async def send_prompt(self, model, text): - messages = [{ "role": "user", "content": text, "options": { "num_gpu": 100, "main_gpu": 0 } }] - async for part in await self.client.chat(model=model, messages=messages, stream=True): + async def send_prompt(self, model, text="", num_gpu=10): + messages = [{ "role": "user", "content": text }] + GLib.idle_add(self.window.hold_prompt, True) + async for part in await self.client.chat(model=model, messages=messages, stream=True, options={ "num_gpu": num_gpu }): GLib.idle_add(self.window.add_response_part, part.message.content) - GLib.idle_add(self.window.entry.set_sensitive, True) + GLib.idle_add(self.window.unhold_prompt, True) if __name__ == "__main__": app = PyLlama()