fix: only use prompt_template when suffix presents [TAB-46] (#184)

* fix: only use prompt_template when suffix presents

* lint
support-coreml
Meng Zhang 2023-06-03 10:29:04 -07:00 committed by GitHub
parent 2779da3cba
commit 775576b53e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 17 deletions

View File

@ -53,17 +53,3 @@ impl CacheInfo {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_hf() {
let hf_metadata = HFCacheInfo::from("TabbyML/J-350M").await;
assert_eq!(
hf_metadata.transformers_info.auto_model,
"AutoModelForCausalLM"
);
}
}

View File

@ -33,7 +33,7 @@ pub struct Segments {
/// Content that appears after the cursor in the editor window.
#[schema(example = "\n return fib(n - 1) + fib(n - 2)")]
suffix: String,
suffix: Option<String>,
}
#[derive(Serialize, Deserialize, ToSchema, Clone, Debug)]
@ -65,8 +65,13 @@ pub async fn completion(
let prompt = if let Some(Segments { prefix, suffix }) = request.segments {
if let Some(prompt_template) = &state.prompt_template {
strfmt!(prompt_template, prefix => prefix, suffix => suffix)
.expect("Failed to format prompt")
if let Some(suffix) = suffix {
strfmt!(prompt_template, prefix => prefix, suffix => suffix)
.expect("Failed to format prompt")
} else {
// If suffix is empty, just returns prefix.
prefix
}
} else {
// If there's no prompt template, just use prefix.
prefix