-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add tokenize & detokenize to client, fix typos
- implemented client code for the `/tokenize` & `/detokenize` endpoints - added docstring examples
- Loading branch information
1 parent
ececcef
commit e2b37f8
Showing
8 changed files
with
296 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,57 @@ | ||
use crate::Task; | ||
use serde::{Deserialize, Serialize}; | ||
|
||
/// Input for a [crate::Client::detokenize] request. | ||
pub struct TaskDetokenization<'a> { | ||
/// List of token ids which should be detokenized into text. | ||
pub token_ids: &'a [u32], | ||
} | ||
|
||
/// Body send to the Aleph Alpha API on the POST `/detokenize` route | ||
#[derive(Serialize, Debug)] | ||
struct BodyDetokenization<'a> { | ||
/// Name of the model tasked with completing the prompt. E.g. `luminous-base"`. | ||
pub model: &'a str, | ||
/// List of ids to detokenize. | ||
pub token_ids: &'a [u32], | ||
} | ||
|
||
#[derive(Deserialize, Debug, PartialEq, Eq)] | ||
pub struct ResponseDetokenization { | ||
pub result: String, | ||
} | ||
|
||
#[derive(Debug, PartialEq, Eq)] | ||
pub struct DetokenizationOutput { | ||
pub result: String, | ||
} | ||
|
||
impl From<ResponseDetokenization> for DetokenizationOutput { | ||
fn from(response: ResponseDetokenization) -> Self { | ||
Self { | ||
result: response.result, | ||
} | ||
} | ||
} | ||
|
||
impl<'a> Task for TaskDetokenization<'a> { | ||
type Output = DetokenizationOutput; | ||
type ResponseBody = ResponseDetokenization; | ||
|
||
fn build_request( | ||
&self, | ||
client: &reqwest::Client, | ||
base: &str, | ||
model: &str, | ||
) -> reqwest::RequestBuilder { | ||
let body = BodyDetokenization { | ||
model, | ||
token_ids: &self.token_ids, | ||
}; | ||
client.post(format!("{base}/detokenize")).json(&body) | ||
} | ||
|
||
fn body_to_output(&self, response: Self::ResponseBody) -> Self::Output { | ||
DetokenizationOutput::from(response) | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,91 @@ | ||
use crate::Task; | ||
use serde::{Deserialize, Serialize}; | ||
|
||
/// Input for a [crate::Client::tokenize] request. | ||
pub struct TaskTokenization<'a> { | ||
/// The text prompt which should be converted into tokens | ||
pub prompt: &'a str, | ||
|
||
/// Specify `true` to return text-tokens. | ||
pub tokens: bool, | ||
|
||
/// Specify `true` to return numeric token-ids. | ||
pub token_ids: bool, | ||
} | ||
|
||
impl<'a> From<&'a str> for TaskTokenization<'a> { | ||
fn from(prompt: &'a str) -> TaskTokenization { | ||
TaskTokenization { | ||
prompt, | ||
tokens: true, | ||
token_ids: true, | ||
} | ||
} | ||
} | ||
|
||
impl TaskTokenization<'_> { | ||
pub fn new(prompt: &str, tokens: bool, token_ids: bool) -> TaskTokenization { | ||
TaskTokenization { | ||
prompt, | ||
tokens, | ||
token_ids, | ||
} | ||
} | ||
} | ||
|
||
#[derive(Serialize, Debug)] | ||
struct BodyTokenization<'a> { | ||
/// Name of the model tasked with completing the prompt. E.g. `luminous-base`. | ||
pub model: &'a str, | ||
/// String to tokenize. | ||
pub prompt: &'a str, | ||
/// Set this value to `true` to return text-tokens. | ||
pub tokens: bool, | ||
/// Set this value to `true` to return numeric token-ids. | ||
pub token_ids: bool, | ||
} | ||
|
||
#[derive(Deserialize, Debug, PartialEq, Eq)] | ||
pub struct ResponseTokenization { | ||
pub tokens: Option<Vec<String>>, | ||
pub token_ids: Option<Vec<u32>>, | ||
} | ||
|
||
#[derive(Debug, PartialEq)] | ||
pub struct TokenizationOutput { | ||
pub tokens: Option<Vec<String>>, | ||
pub token_ids: Option<Vec<u32>>, | ||
} | ||
|
||
impl From<ResponseTokenization> for TokenizationOutput { | ||
fn from(response: ResponseTokenization) -> Self { | ||
Self { | ||
tokens: response.tokens, | ||
token_ids: response.token_ids, | ||
} | ||
} | ||
} | ||
|
||
impl Task for TaskTokenization<'_> { | ||
type Output = TokenizationOutput; | ||
type ResponseBody = ResponseTokenization; | ||
|
||
fn build_request( | ||
&self, | ||
client: &reqwest::Client, | ||
base: &str, | ||
model: &str, | ||
) -> reqwest::RequestBuilder { | ||
let body = BodyTokenization { | ||
model, | ||
prompt: &self.prompt, | ||
tokens: self.tokens, | ||
token_ids: self.token_ids, | ||
}; | ||
client.post(format!("{base}/tokenize")).json(&body) | ||
} | ||
|
||
fn body_to_output(&self, response: Self::ResponseBody) -> Self::Output { | ||
TokenizationOutput::from(response) | ||
} | ||
} |
Oops, something went wrong.