Merge pull request #4 from jinshen1983/main

Add optional parameter model to tokenize/detokenize

Author: andreibondarev

Gemfile.lock

@@ -71,6 +71,7 @@ GEM
 
 PLATFORMS
   x86_64-darwin-19
+  x86_64-darwin-21
   x86_64-linux
 
 DEPENDENCIES

lib/cohere/client.rb

@@ -83,16 +83,16 @@ def classify(
       response.body
     end
 
-    def tokenize(text:)
+    def tokenize(text:, model: nil)
       response = connection.post("tokenize") do |req|
-        req.body = {text: text}
+        req.body = model.nil? ? {text: text} : {text: text, model: model}
       end
       response.body
     end
 
-    def detokenize(tokens:)
+    def detokenize(tokens:, model: nil)
       response = connection.post("detokenize") do |req|
-        req.body = {tokens: tokens}
+        req.body = model.nil? ? {tokens: tokens} : {tokens: tokens, model: model}
       end
       response.body
     end

spec/cohere/client_spec.rb

@@ -91,6 +91,24 @@
     end
   end
 
+  describe "#tokenize_with_model" do
+    let(:tokenize_result) { JSON.parse(File.read("spec/fixtures/tokenize_result.json")) }
+    let(:response) { OpenStruct.new(body: tokenize_result) }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+        .with("tokenize")
+        .and_return(response)
+    end
+
+    it "returns a response" do
+      expect(instance.tokenize(
+        text: "Hello, world!",
+        model: "base"
+      ).dig("tokens")).to eq([33555, 1114, 34])
+    end
+  end
+
   describe "#detokenize" do
     let(:detokenize_result) { JSON.parse(File.read("spec/fixtures/detokenize_result.json")) }
     let(:response) { OpenStruct.new(body: detokenize_result) }
@@ -108,6 +126,24 @@
     end
   end
 
+  describe "#detokenize_with_model" do
+    let(:detokenize_result) { JSON.parse(File.read("spec/fixtures/detokenize_result.json")) }
+    let(:response) { OpenStruct.new(body: detokenize_result) }
+
+    before do
+      allow_any_instance_of(Faraday::Connection).to receive(:post)
+        .with("detokenize")
+        .and_return(response)
+    end
+
+    it "returns a response" do
+      expect(instance.detokenize(
+        tokens: [33555, 1114, 34],
+        model: "base"
+      ).dig("text")).to eq("hello world!")
+    end
+  end
+
   describe "#detect_language" do
     let(:detect_language_result) { JSON.parse(File.read("spec/fixtures/detect_language_result.json")) }
     let(:response) { OpenStruct.new(body: detect_language_result) }