Make the app nicer!

Dockerfile

@@ -5,7 +5,8 @@ WORKDIR /code
 # Install stable packages from CRAN
 RUN install2.r --error \
     ggExtra \
-    shiny
+    shiny \
+    callr
 
 # Install Rust for tok
 

app.R

@@ -1,53 +1,125 @@
 library(shiny)
 library(bslib)
-library(dplyr)
-library(ggplot2)
 library(minhub)
+source("model-session.R")
 
-model <- gptneox()
+repo <- "stabilityai/stablelm-tuned-alpha-3b"
+sess <- model_session$new()
+model_loaded <- sess$load_model(repo)
 
-# Find subset of columns that are suitable for scatter plot
-df_num <- df |> select(where(is.numeric), -Year)
+max_n_tokens <- 100
+system_prompt = "<|SYSTEM|># StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+"
 
-ui <- page_fillable(theme = bs_theme(bootswatch = "minty"),
-  layout_sidebar(fillable = TRUE,
-    sidebar(
-      varSelectInput("xvar", "X variable", df_num, selected = "Bill Length (mm)"),
-      varSelectInput("yvar", "Y variable", df_num, selected = "Bill Depth (mm)"),
-      checkboxGroupInput("species", "Filter by species",
-        choices = unique(df$Species), selected = unique(df$Species)
-      ),
-      hr(), # Add a horizontal rule
-      checkboxInput("by_species", "Show species", TRUE),
-      checkboxInput("show_margins", "Show marginal plots", TRUE),
-      checkboxInput("smooth", "Add smoother"),
-    ),
-    plotOutput("scatter")
+ui <- page_fillable(
+  theme = bs_theme(bootswatch = "minty"),
+  shinyjs::useShinyjs(),
+  card(
+    height="90%",
+    heights_equal = "row",
+    width = 1,
+    fillable = FALSE,
+    card_body(id = "messages", gap = 5, fillable = FALSE)
+  ),
+  layout_column_wrap(
+    width = 1/2,
+    textInput("prompt", label = NULL, width="100%"),
+    actionButton("send", "Loading model...", width = "100%")
   )
 )
 
 server <- function(input, output, session) {
-  subsetted <- reactive({
-    req(input$species)
-    df |> filter(Species %in% input$species)
-  })
   
-  output$scatter <- renderPlot({
-    p <- ggplot(subsetted(), aes(!!input$xvar, !!input$yvar)) + list(
-      theme(legend.position = "bottom"),
-      if (input$by_species) aes(color=Species),
-      geom_point(),
-      if (input$smooth) geom_smooth()
-    )
-
-    if (input$show_margins) {
-      margin_type <- if (input$by_species) "density" else "histogram"
-      p <- p |> ggExtra::ggMarginal(type = margin_type, margins = "both",
-        size = 8, groupColour = input$by_species, groupFill = input$by_species)
+  prompt <- reactiveVal(value = system_prompt)
+  n_tokens <- reactiveVal(value = 0)
+  
+  observeEvent(input$send, {
+    if (is.null(input$prompt) || input$prompt == "") {
+      return()
     }
+    shinyjs::disable("send")
+    updateActionButton(inputId = "send", label = "Waiting for model")
+    insert_message(as.character(glue::glue("🤗: {input$prompt}")))  
+    
+    # we modify the prompt to trigger the 'next_token' reactive
+    prompt(paste0(prompt(), "<|USER|>", input$prompt, "<|ASSISTANT|>")) 
+  })
+  
+  next_token <- eventReactive(prompt(), ignoreInit = TRUE, {
+    prompt() %>% 
+      sess$generate()
+  })
+  
+  observeEvent(next_token(), {
+    tok <- next_token()
+    n_tokens(n_tokens() + 1)
+    
+    tok %>% promises::then(function(tok) {
+      if (n_tokens() == 1) {
+        insert_message(paste0("🤖: ", tok), append = FALSE)
+      } else {
+        insert_message(tok, append = TRUE)
+      }
+      
+      if (tok != "" && n_tokens() < max_n_tokens) {
+        prompt(paste0(prompt(), tok))
+      } else {
+        shinyjs::enable("send")
+        updateActionButton(inputId = "send", label = "Send")
+        n_tokens(0)
+      }
+    })
+  })
+  
+  # we need this observer to make sure that during the event loop the
+  # tasks are resolved.
+  observe({
+    invalidateLater(5000, session)
+    sess$sess$poll_process(1)
+  })
+  
+  # Observer used at app startup time to allow using the 'Send' button once the
+  # model has been loaded.
+  observe({
+    ready <- sess$sess$poll_process(1) == "ready"
+    send <- isolate(input$send)
     
-    p
-  }, res = 100)
+    if (send == 0 && !ready) {
+      invalidateLater(1000, session)  
+    }
+  
+    if (ready) {
+      shinyjs::enable("send")
+      updateActionButton(inputId = "send", label = "Send")
+    } else {
+      shinyjs::disable("send")
+    }
+  })
 }
 
+message_id <- 0
+insert_message <- function(msg, append = FALSE) {
+  if (!append) {
+    id <- message_id <<- message_id + 1
+    insertUI(
+      "#messages", 
+      "beforeEnd", 
+      immediate = TRUE,
+      ui = card(card_body(p(id = paste0("msg-",id), msg)), style="margin-bottom:5px;")
+    )
+  } else {
+    id <- message_id
+    shinyjs::runjs(glue::glue(
+      "document.getElementById('msg-{id}').textContent += '{msg}'"
+    ))
+  }
+  # scroll to bottom
+  shinyjs::runjs("var elem = document.getElementById('messages'); elem.scrollTop = elem.scrollHeight;")
+}
+
+
 shinyApp(ui, server)

gptneox.Rproj

@@ -0,0 +1,13 @@
+Version: 1.0
+
+RestoreWorkspace: Default
+SaveWorkspace: Default
+AlwaysSaveHistory: Default
+
+EnableCodeIndexing: Yes
+UseSpacesForTab: Yes
+NumSpacesForTab: 2
+Encoding: UTF-8
+
+RnwWeave: knitr
+LaTeX: pdfLaTeX

model-session.R

@@ -0,0 +1,45 @@
+source("promise-session.R")
+
+# A wrapper a around the promise session that controls model loading and
+# querying given a prompt
+model_session <- R6::R6Class(
+  lock_objects = FALSE,
+  public = list(
+    initialize = function() {
+      self$sess <- promise_session$new()
+      self$temperature <- 1
+      self$top_k <- 50
+    },
+    load_model = function(repo) {
+      self$sess$call(args = list(repo = repo), function(repo) {
+        library(torch)
+        library(zeallot)
+        library(minhub)
+        model <<- minhub::gptneox_from_pretrained(repo)
+        model$eval()
+        model$to(dtype = torch_float())
+        tok <<- tok::tokenizer$from_pretrained(repo)
+        "done"
+      })
+    },
+    generate = function(prompt) {
+      args <- list(
+        prompt = prompt, 
+        temperature = self$temperature,
+        top_k = self$top_k
+      )
+      self$sess$call(args = args, function(prompt, temperature, top_k) {
+        idx <- torch_tensor(tok$encode(prompt)$ids)$view(c(1, -1))
+        with_no_grad({
+          logits <- model(idx + 1L)
+        })
+        logits <- logits[,-1,]/temperature
+        c(prob, ind) %<-% logits$topk(top_k)
+        logits <- torch_full_like(logits, -Inf)$scatter_(-1, ind, prob)
+        logits <- nnf_softmax(logits, dim = -1)
+        id_next <- torch::torch_multinomial(logits, num_samples = 1) - 1L
+        tok$decode(as.integer(id_next))  
+      })
+    }
+  )
+)

promise-session.R

@@ -0,0 +1,67 @@
+
+# Small utility class that wraps a `callr::r_session` to return promises when
+# executing `sess$call()`.
+# Only one promise is resolve per time in fifo way.
+promise_session <- R6::R6Class(
+  lock_objects = FALSE,
+  public = list(
+    initialize = function() {
+      self$sess <- callr::r_session$new()
+      self$is_running <- FALSE
+    },
+    call = function(func, args = list()) {
+      self$poll_process()
+      promises::promise(function(resolve, reject) {
+        self$push_task(func, args, resolve, reject)
+        later::later(self$poll_process, 1)
+      })
+    },
+    push_task = function(func, args, resolve, reject) {
+      self$tasks[[length(self$tasks) + 1]] <- list(
+        func = func, 
+        args = args, 
+        resolve = resolve, 
+        reject = reject
+      )
+      cat("task pushed, now we have ", length(self$tasks), " on queue\n")
+      self$run_task()
+      invisible(NULL)
+    },
+    run_task = function() {
+      if (self$is_running) return(NULL)
+      if (length(self$tasks) == 0) return(NULL)
+      
+      self$is_running <- TRUE
+      task <- self$tasks[[1]]
+      self$sess$call(task$func, args = task$args)
+    },
+    resolve_task = function() {
+      out <- self$sess$read()
+      if (!is.null(out$error)) {
+        self$tasks[[1]]$reject(out$error)
+      } else {
+        self$tasks[[1]]$resolve(out$result)
+      }
+      
+      self$tasks <- self$tasks[-1]
+      self$is_running <- FALSE
+      
+      self$run_task()
+    },
+    poll_process = function(timeout = 1) {
+      if (!self$is_running) return("ready")
+      poll_state <- self$sess$poll_process(timeout)
+      if (poll_state == "ready") {
+        self$resolve_task()
+      }
+      poll_state
+    }
+  )
+)
+
+# sess <- promise_session$new()
+# f <- sess$call(function(a) {
+#   10 + 1
+# }, list(1))
+# sess$poll_process()
+