use std::collections::HashMap;
use serde::{Deserialize, Serialize};
#[derive(Serialize, Debug, Clone)]
pub struct CompletionArguments {
pub model: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub suffix: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub n: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) stream: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub logprobs: Option<u8>,
#[serde(skip_serializing_if = "Option::is_none")]
pub echo: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub presence_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub frequency_penalty: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub best_of: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
}
impl CompletionArguments {
pub fn new(model: impl AsRef<str>, prompt: String) -> CompletionArguments {
CompletionArguments {
model: model.as_ref().to_owned(),
prompt: Some(prompt),
suffix: None,
max_tokens: None,
temperature: None,
top_p: None,
n: None,
stream: None,
logprobs: None,
echo: None,
stop: None,
presence_penalty: None,
frequency_penalty: None,
best_of: None,
user: None,
}
}
}
#[derive(Deserialize, Debug, Clone)]
pub struct CompletionResponse {
pub id: String,
pub created: u32,
pub model: String,
pub choices: Vec<Choice>,
pub usage: Usage,
}
impl std::fmt::Display for CompletionResponse {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.choices[0].text)?;
Ok(())
}
}
#[derive(Deserialize, Debug, Clone)]
pub struct Choice {
pub text: String,
pub index: u32,
pub logprobs: Option<LogProbs>,
pub finish_reason: String,
}
#[derive(Deserialize, Debug, Clone)]
pub struct LogProbs {
pub tokens: Vec<String>,
pub token_logprobs: Vec<f32>,
pub top_logprobs: Vec<HashMap<String, f32>>,
pub text_offset: Vec<u32>,
}
#[derive(Deserialize, Debug, Clone)]
pub struct Usage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
pub total_tokens: u32,
}