Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Load language models in the background

Browse source
This commit is contained in:
Antonio Scandurra 2024-01-24 13:36:44 +01:00
parent 92b0184036
commit 22046ef9a7
4 changed files with 83 additions and 69 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/ai/src/providers/open_ai/completion.rs
View file

@ -201,8 +201,10 @@ pub struct OpenAICompletionProvider {
}
impl OpenAICompletionProvider {
pub fn new(model_name: &str, executor: BackgroundExecutor) -> Self {
let model = OpenAILanguageModel::load(model_name);
pub async fn new(model_name: String, executor: BackgroundExecutor) -> Self {
let model = executor
.spawn(async move { OpenAILanguageModel::load(&model_name) })
.await;
let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
Self {
model,

7
crates/ai/src/providers/open_ai/embedding.rs
View file

@ -67,11 +67,14 @@ struct OpenAIEmbeddingUsage {
}
impl OpenAIEmbeddingProvider {
pub fn new(client: Arc<dyn HttpClient>, executor: BackgroundExecutor) -> Self {
pub async fn new(client: Arc<dyn HttpClient>, executor: BackgroundExecutor) -> Self {
let (rate_limit_count_tx, rate_limit_count_rx) = watch::channel_with(None);
let rate_limit_count_tx = Arc::new(Mutex::new(rate_limit_count_tx));
let model = OpenAILanguageModel::load("text-embedding-ada-002");
// Loading the model is expensive, so ensure this runs off the main thread.
let model = executor
.spawn(async move { OpenAILanguageModel::load("text-embedding-ada-002") })
.await;
let credential = Arc::new(RwLock::new(ProviderCredential::NoCredentials));
OpenAIEmbeddingProvider {