generate code refactoring progress 1/3

src/models/minicpm4/generate.rs

@@ -1,25 +1,18 @@
-use crate::models::common::MultiModalData;
-use crate::models::common::generate::{
-    GenerationContext, generate_generic, generate_stream_generic,
-};
-use crate::params::chat::{
-    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
-};
+use crate::models::common::generate::{GenerationDataProvider, PrepareData};
+
 use anyhow::Result;
 use candle_core::{DType, Device};
 use candle_nn::VarBuilder;
-use rocket::futures::Stream;
 
 use crate::models::minicpm4::config::MiniCPM4Config;
 use crate::models::minicpm4::model::MiniCPMModel;
-// use crate::models::GenerateStream;
 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 MiniCPM4GenerateModel<'a> {
     chat_template: ChatTemplate<'a>,
     tokenizer: TokenizerModel,
-    minicpm: MiniCPMModel,
+    model: MiniCPMModel,
     device: Device,
     model_name: String,
 }
@@ -35,7 +28,7 @@ impl<'a> MiniCPM4GenerateModel<'a> {
         let dtype = get_dtype(dtype, cfg_dtype);
         let model_list = find_type_files(path, "safetensors")?;
         let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
-        let minicpm = MiniCPMModel::new(vb, cfg)?;
+        let model = MiniCPMModel::new(vb, cfg)?;
         let model_name = std::path::Path::new(path)
             .file_name()
             .and_then(|s| s.to_str())
@@ -44,73 +37,25 @@ impl<'a> MiniCPM4GenerateModel<'a> {
         Ok(MiniCPM4GenerateModel {
             chat_template,
             tokenizer,
-            minicpm,
+            model,
             device: device.clone(),
             model_name,
         })
     }
 }
 
-impl<'a> GenerateModel for MiniCPM4GenerateModel<'a> {
-    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
-        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+impl<'a> GenerationDataProvider for MiniCPM4GenerateModel<'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.minicpm,
-            &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.minicpm,
-            &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!(MiniCPM4GenerateModel<'a>);