remove fetch method from ModelTrait

atoma-network · Mar 31, 2024 · f673dea · f673dea
1 parent cddb534
commit f673dea
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 10 deletions.
diff --git a/atoma-inference/src/model_thread.rs b/atoma-inference/src/model_thread.rs
@@ -113,7 +113,7 @@ where
     ) -> Result<(Self, Vec<ModelThreadHandle<T, U>>), ModelThreadError>
     where
         F: ApiTrait,
-        M: ModelTrait<FetchApi = F, Input = T::ModelInput, Output = U::ModelOutput>
+        M: ModelTrait<Input = T::ModelInput, Output = U::ModelOutput>
             + Send
             + 'static,
     {

diff --git a/atoma-inference/src/models/mod.rs b/atoma-inference/src/models/mod.rs
@@ -1,6 +1,5 @@
 use std::path::PathBuf;
 
-use crate::{apis::ApiTrait, models::config::ModelConfig};
 use ed25519_consensus::VerificationKey as PublicKey;
 use thiserror::Error;
 
@@ -16,11 +15,9 @@ pub trait ModelBuilder {
 
 pub trait ModelTrait {
     type Builder: Send + Sync + 'static;
-    type FetchApi: ApiTrait + Send + Sync + 'static;
     type Input;
     type Output;
 
-    fn fetch(api: &Self::FetchApi, config: ModelConfig) -> Result<(), ModelError>;
     fn load(filenames: Vec<PathBuf>) -> Result<Self, ModelError>
     where
         Self: Sized;

diff --git a/atoma-inference/src/service.rs b/atoma-inference/src/service.rs
@@ -35,7 +35,7 @@ where
         request_receiver: Receiver<T>,
     ) -> Result<Self, ModelServiceError>
     where
-        M: ModelTrait<FetchApi = F, Input = T::ModelInput, Output = U::ModelOutput>
+        M: ModelTrait<Input = T::ModelInput, Output = U::ModelOutput>
             + Send
             + 'static,
         F: ApiTrait,
@@ -200,14 +200,9 @@ mod tests {
 
     impl ModelTrait for TestModelInstance {
         type Builder = ();
-        type FetchApi = MockApi;
         type Input = ();
         type Output = ();
 
-        fn fetch(_: &Self::FetchApi, _: ModelConfig) -> Result<(), crate::models::ModelError> {
-            Ok(())
-        }
-
         fn load(_: Vec<PathBuf>) -> Result<Self, crate::models::ModelError> {
             Ok(Self {})
         }