search-with-llama.el

;;; search-with-llama.el --- Find an answer to your query using llama!  -*- lexical-binding: t; -*-

;; Copyright (C) 2023  Kiran Gopinathan

;; Author: Kiran Gopinathan
;; Keywords: 

;; This program is free software; you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.

;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.

;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.


;;; Code:


(defcustom swl-llama-binary "/home/kirang/Documents/code/devils-code-ml/llama.cpp/main"
  "Path to the llama.cpp compiled binary")

(defcustom swl-llama-model-path "/home/kirang/Documents/code/devils-code-ml/llama-models/"
  "path to LLama-models")

(defcustom swl-llama-args '("-n" "512")
  "Arguments to pass to lama")


(defvar swl-ggml-f16-file-name "ggml-model-f16.bin")
(defvar swl-ggml-quantized-file-name "ggml-model-q4_0.bin")

(defvar swl-model-sizes '("7B" "13B" "30B" "60B"))

(defvar swl-current-process-buffer nil "Mini-buffer currently being used to display the process.")

(defvar swl-query-hist nil)

(defvar-local swl-process-state nil "State of the swl buffer.")
(defvar-local swl-process-start-point nil "Point at which query starts.")
(defvar-local swl-process-cmd nil "Original query command running in the buffer.")

(defvar swl-result-mode-map
  (let ((map (make-sparse-keymap)))
    (suppress-keymap map)
    (define-key map "q" 'swl-quit)
    (define-key map "r" 'swl-restart-query)
    map))

(defun swl-make-buffer-command (model command &optional full)
  "Return a string that invokes llama.cpp with a query COMMAND and model MODEL."
  (let ((model-file-name (if full swl-ggml-f16-file-name swl-ggml-quantized-file-name)))
    (format "%s -m %s%s/%s -p \"%s:\" %s"
            swl-llama-binary
            (file-name-as-directory swl-llama-model-path) model
            model-file-name
            command
            (string-join swl-llama-args " "))))

(defun swl-clear-running-process ()
  (when (and swl-current-process-buffer (buffer-live-p swl-current-process-buffer))
    (let  ((process (get-buffer-process swl-current-process-buffer)))
      (if process (kill-process process)))
    (with-current-buffer swl-current-process-buffer 
      (erase-buffer))))

(defun swl-quit ()
  (interactive)
  (swl-clear-running-process)
  (when swl-current-process-buffer
    (with-current-buffer swl-current-process-buffer
      (setq swl-process-cmd nil)
      (setq swl-process-state nil)))
  (quit-window))

(defun swl-run-query (cmd)
  (if swl-current-process-buffer
      (swl-clear-running-process))
  (with-current-buffer swl-current-process-buffer
    (setq swl-process-state nil)
    (setq swl-process-cmd cmd)
    (goto-char 0)
    (insert "CMD: ") (insert cmd) (insert "\n")
    (insert "Press r to re-run query and q to quit.\n")
    (goto-char (point-max))
    (setq swl-process-start-point (make-marker))
    (move-marker swl-process-start-point (point))
    (use-local-map swl-result-mode-map))
  (let (proc)
    (setq proc (start-process-shell-command
                "swl-search-process"
                swl-current-process-buffer
                cmd))
    (when (and proc (processp proc))
      (set-process-filter proc #'swl-filter))))

(defun swl-restart-query ()
  (interactive)
  (swl-clear-running-process)
  (when (and swl-current-process-buffer (buffer-live-p swl-current-process-buffer))
    (with-current-buffer swl-current-process-buffer
      (when swl-process-cmd
        (swl-run-query swl-process-cmd)))))

(defun swl-filter (process event)
  (when (and swl-current-process-buffer (buffer-live-p swl-current-process-buffer) (process-live-p process))
    (with-current-buffer swl-current-process-buffer
      (save-excursion
        (goto-char (process-mark process))
        (insert event)
        (set-marker (process-mark process) (point))
        (unless swl-process-state
          (goto-char (point-min))
          (setq swl-process-state
                (search-forward-regexp "sampling parameters:.*\n" nil t))
          (if swl-process-state
              (delete-region (marker-position swl-process-start-point) swl-process-state)))))))

(defun search-with-llama (query &optional model-size model-type)
  (interactive
   (list
    (completing-read "Find a guide on: " swl-query-hist nil nil nil 'swl-query-hist)
    (if current-prefix-arg
        (completing-read "Model: " swl-model-sizes nil t "7B")
      "7B")
    (if current-prefix-arg
        (completing-read "Model type: " '(quantized f16) nil t)
      "quantized")))
  (unless (and swl-current-process-buffer (buffer-live-p swl-current-process-buffer))
    (setq swl-current-process-buffer (get-buffer-create "*swl-process-buffer*")))
  (let (cmd)
    (setq cmd (swl-make-buffer-command model-size query
                                         (and model-type
                                              (stringp model-type)
                                              (string-equal model-type "f16"))))
    (swl-run-query cmd)
    (display-buffer-at-bottom swl-current-process-buffer '(previous-window))
    (pop-to-buffer swl-current-process-buffer)))


(provide 'search-with-llama)
;;; search-with-llama.el ends here