generate code refactoring progress 1/3

src/models/minicpm5/generate.rs

@@ -1,19 +1,13 @@
-use crate::models::common::MultiModalData;
-use crate::models::common::generate::{
-    GenerationContext, generate_generic, generate_stream_generic,
-};
+use crate::models::common::generate::{GenerationDataProvider, PrepareData};
 use crate::models::llama::LlamaForCausalLM;
 use crate::models::minicpm5::config::MiniCPM5Config;
-use crate::params::chat::{
-    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
-};
+
 use anyhow::Result;
 use candle_core::{DType, Device};
 use candle_nn::VarBuilder;
-use rocket::futures::Stream;
 
 use crate::utils::{find_type_files, get_device, get_dtype};
-use crate::{chat_template::ChatTemplate, models::GenerateModel, tokenizer::TokenizerModel};
+use crate::{chat_template::ChatTemplate, tokenizer::TokenizerModel};
 
 pub struct MiniCPM5GenerateModel<'a> {
     chat_template: ChatTemplate<'a>,
@@ -70,66 +64,18 @@ impl<'a> MiniCPM5GenerateModel<'a> {
     }
 }
 
-impl<'a> GenerateModel for MiniCPM5GenerateModel<'a> {
-    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
-        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+impl<'a> GenerationDataProvider for MiniCPM5GenerateModel<'a> {
+    fn get_data(&self, mes: &crate::params::chat::ChatCompletionParameters) -> Result<PrepareData> {
+        let mes_render = self.chat_template.apply_chat_template(mes)?;
+        let in_reasoning = self.is_in_reasoning(&mes_render);
         let input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
-        let seed = mes.seed.unwrap_or(34562) as u64;
-        let sample_len = mes.max_tokens.unwrap_or(2048);
-        let mut ctx = GenerationContext::new(
-            mes.temperature,
-            mes.top_p,
-            None,
-            mes.repeat_penalty,
-            mes.repeat_last_n,
-            seed,
-            input_ids.dim(1)?,
-            sample_len,
-            self.device.clone(),
-        );
-
-        let data = MultiModalData::new(vec![]);
-        generate_generic(
-            &mut self.model,
-            &self.tokenizer,
+        let multi_model_data = self.get_multi_model_data();
+        Ok(PrepareData {
+            in_reasoning,
             input_ids,
-            data,
-            &mut ctx,
-            &self.model_name,
-        )
-    }
-    fn generate_stream(
-        &mut self,
-        mes: ChatCompletionParameters,
-    ) -> Result<
-        Box<
-            dyn Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>
-                + Send
-                + Unpin
-                + '_,
-        >,
-    > {
-        let seed = mes.seed.unwrap_or(34562) as u64;
-        let mes_render = self.chat_template.apply_chat_template(&mes)?;
-        let input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
-        let data = MultiModalData::new(vec![]);
-        let sample_len = mes.max_tokens.unwrap_or(512);
-        let stream = generate_stream_generic(
-            &mut self.model,
-            &self.tokenizer,
-            input_ids,
-            data,
-            mes.temperature,
-            mes.top_p,
-            None,
-            mes.repeat_penalty,
-            mes.repeat_last_n,
-            seed,
-            sample_len,
-            false,
-            &self.device,
-            &self.model_name,
-        )?;
-        Ok(Box::new(Box::pin(stream)))
+            multi_model_data,
+        })
     }
 }
+
+crate::impl_generate_model!(MiniCPM5GenerateModel<'a>);