Adds classify Ollama tests, adds table tests to sentiment

mlverse · Sep 12, 2024 · b5e9fce · b5e9fce
1 parent 888cd18
commit b5e9fce
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 2 deletions.
diff --git a/tests/testthat/_snaps/llm-classify.md b/tests/testthat/_snaps/llm-classify.md
@@ -0,0 +1,44 @@
+# Classify on Ollama works
+
+    Code
+      llm_classify(reviews, review, labels = c("appliance", "computer"))
+    Output
+                                                                                    review
+      1                 This has been the best TV I've ever used. Great screen, and sound.
+      2          I regret buying this laptop. It is too slow and the keyboard is too noisy
+      3 Not sure how to feel about my new washing machine. Great color, but hard to figure
+        .classify
+      1 appliance
+      2  computer
+      3 appliance
+
+---
+
+    Code
+      llm_classify(reviews, review, pred_name = "new", labels = c("appliance",
+        "computer"))
+    Output
+                                                                                    review
+      1                 This has been the best TV I've ever used. Great screen, and sound.
+      2          I regret buying this laptop. It is too slow and the keyboard is too noisy
+      3 Not sure how to feel about my new washing machine. Great color, but hard to figure
+              new
+      1 appliance
+      2  computer
+      3 appliance
+
+---
+
+    Code
+      llm_classify(reviews, review, pred_name = "new", labels = c("appliance",
+        "computer"), additional_prompt = "Consider all laptops as appliances.")
+    Output
+                                                                                    review
+      1                 This has been the best TV I've ever used. Great screen, and sound.
+      2          I regret buying this laptop. It is too slow and the keyboard is too noisy
+      3 Not sure how to feel about my new washing machine. Great color, but hard to figure
+              new
+      1 appliance
+      2 appliance
+      3 appliance
+
diff --git a/tests/testthat/_snaps/llm-sentiment.md b/tests/testthat/_snaps/llm-sentiment.md
@@ -20,3 +20,31 @@
     Output
       [1] "positive" "negative" "positive"
 
+---
+
+    Code
+      llm_sentiment(reviews, review)
+    Output
+                                                                                    review
+      1                 This has been the best TV I've ever used. Great screen, and sound.
+      2          I regret buying this laptop. It is too slow and the keyboard is too noisy
+      3 Not sure how to feel about my new washing machine. Great color, but hard to figure
+        .sentiment
+      1   positive
+      2   negative
+      3    neutral
+
+---
+
+    Code
+      llm_sentiment(reviews, review, pred_name = "new")
+    Output
+                                                                                    review
+      1                 This has been the best TV I've ever used. Great screen, and sound.
+      2          I regret buying this laptop. It is too slow and the keyboard is too noisy
+      3 Not sure how to feel about my new washing machine. Great color, but hard to figure
+             new
+      1 positive
+      2 negative
+      3  neutral
+
diff --git a/tests/testthat/helper-ollama.R b/tests/testthat/helper-ollama.R
@@ -2,6 +2,8 @@ skip_if_no_ollama <- function() {
   con <- ollamar::test_connection()
   if (con$status_code != 200) {
     skip("No Ollama found")
+  } else {
+    llm_use("ollama", "llama3.1", seed = 100, .silent = TRUE)  
   }
 }
 
@@ -14,5 +16,5 @@ reviews_vec <- function() {
 }
 
 reviews_table <- function() {
-  data.frame(reviews = reviews_vec())
+  data.frame(review = reviews_vec())
 }
diff --git a/tests/testthat/test-llm-classify.R b/tests/testthat/test-llm-classify.R
@@ -25,3 +25,32 @@ test_that("Classify works", {
     data.frame(x = test_text, new = test_text)
   )
 })
+
+test_that("Classify on Ollama works", {
+  skip_if_no_ollama()
+  reviews <- reviews_table()
+  expect_snapshot(
+    llm_classify(
+      reviews, 
+      review,
+      labels = c("appliance", "computer")
+      )
+    )
+  expect_snapshot(
+    llm_classify(
+      reviews, 
+      review, 
+      pred_name = "new",
+      labels = c("appliance", "computer")
+      )
+    )
+  expect_snapshot(
+    llm_classify(
+      reviews, 
+      review, 
+      pred_name = "new", 
+      labels = c("appliance", "computer"), 
+      additional_prompt = "Consider all laptops as appliances."
+      )
+    )
+})
diff --git a/tests/testthat/test-llm-sentiment.R b/tests/testthat/test-llm-sentiment.R
@@ -23,7 +23,7 @@ test_that("Sentiment works", {
 test_that("Sentiment on Ollama works", {
   skip_if_no_ollama()
   vec_reviews <- reviews_vec()
-  llm_use("ollama", "llama3.1", seed = 100, .silent = TRUE)
+  reviews <- reviews_table()
   expect_snapshot(llm_vec_sentiment(vec_reviews))
   expect_snapshot(
     llm_vec_sentiment(
@@ -38,4 +38,6 @@ test_that("Sentiment on Ollama works", {
       additional_prompt = "Consider someone not sure as a positive comment."
     )
   )
+  expect_snapshot(llm_sentiment(reviews, review))
+  expect_snapshot(llm_sentiment(reviews, review, pred_name = "new"))
 })