-
Notifications
You must be signed in to change notification settings - Fork 8
/
client.rb
143 lines (130 loc) · 3.83 KB
/
client.rb
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
# frozen_string_literal: true
require "faraday"
require "faraday_middleware"
module Cohere
class Client
attr_reader :api_key, :connection
ENDPOINT_URL = "https://api.cohere.ai/v1"
def initialize(api_key:)
@api_key = api_key
end
# This endpoint generates realistic text conditioned on a given input.
def generate(
prompt:,
model: nil,
num_generations: nil,
max_tokens: nil,
preset: nil,
temperature: nil,
k: nil,
p: nil,
frequency_penalty: nil,
presence_penalty: nil,
end_sequences: nil,
stop_sequences: nil,
return_likelihoods: nil,
logit_bias: nil,
truncate: nil
)
response = connection.post("generate") do |req|
req.body = {prompt: prompt}
req.body[:model] = model if model
req.body[:num_generations] = num_generations if num_generations
req.body[:max_tokens] = max_tokens if max_tokens
req.body[:preset] = preset if preset
req.body[:temperature] = temperature if temperature
req.body[:k] = k if k
req.body[:p] = p if p
req.body[:frequency_penalty] = frequency_penalty if frequency_penalty
req.body[:presence_penalty] = presence_penalty if presence_penalty
req.body[:end_sequences] = end_sequences if end_sequences
req.body[:stop_sequences] = stop_sequences if stop_sequences
req.body[:return_likelihoods] = return_likelihoods if return_likelihoods
req.body[:logit_bias] = logit_bias if logit_bias
req.body[:truncate] = truncate if truncate
end
response.body
end
def embed(
texts:,
model: nil,
truncate: nil
)
response = connection.post("embed") do |req|
req.body = {texts: texts}
req.body[:model] = model if model
req.body[:truncate] = truncate if truncate
end
response.body
end
def classify(
inputs:,
examples:,
model: nil,
present: nil,
truncate: nil
)
response = connection.post("classify") do |req|
req.body = {
inputs: inputs,
examples: examples
}
req.body[:model] = model if model
req.body[:present] = present if present
req.body[:truncate] = truncate if truncate
end
response.body
end
def tokenize(text:)
response = connection.post("tokenize") do |req|
req.body = {text: text}
end
response.body
end
def detokenize(tokens:)
response = connection.post("detokenize") do |req|
req.body = {tokens: tokens}
end
response.body
end
def detect_language(texts:)
response = connection.post("detect-language") do |req|
req.body = {texts: texts}
end
response.body
end
def summarize(
text:,
length: nil,
format: nil,
model: nil,
extractiveness: nil,
temperature: nil,
additional_command: nil
)
response = connection.post("summarize") do |req|
req.body = {text: text}
req.body[:length] = length if length
req.body[:format] = format if format
req.body[:model] = model if model
req.body[:extractiveness] = extractiveness if extractiveness
req.body[:temperature] = temperature if temperature
req.body[:additional_command] = additional_command if additional_command
end
response.body
end
private
# standard:disable Lint/DuplicateMethods
def connection
@connection ||= Faraday.new(url: ENDPOINT_URL) do |faraday|
if api_key
faraday.request :authorization, :Bearer, api_key
end
faraday.request :json
faraday.response :json, content_type: /\bjson$/
faraday.adapter Faraday.default_adapter
end
end
# standard:enable Lint/DuplicateMethods
end
end