diff --git a/.github/workflows/update_r4userembedding.yaml b/.github/workflows/update_r4userembedding.yaml index 4e67b78..d1d079a 100644 --- a/.github/workflows/update_r4userembedding.yaml +++ b/.github/workflows/update_r4userembedding.yaml @@ -1,33 +1,65 @@ -name: Update Userembedding +name: Publish to Dockerhub ( userembedding ) on: - workflow_dispatch: - inputs: - tags: - description: 'Release Tags' + workflow_dispatch: + inputs: + tags: + description: 'Release Tags' jobs: - update_userembedding: - runs-on: ubuntu-latest - steps: - - name: Check out the repo - uses: actions/checkout@v3 - - - name: Set up QEMU - uses: docker/setup-qemu-action@v3 - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 - - - name: Login to GitHub Container Registry - uses: docker/login-action@v2 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_PASS }} - - - name: Build and push Docker image - uses: docker/build-push-action@v3 - with: - push: true - tags: beclab/r4userembedding:${{ github.event.inputs.tags }} - file: Dockerfile.r4userembedding - platforms: linux/amd64,linux/arm64 \ No newline at end of file + publish_dockerhub_amd64: + runs-on: ubuntu-latest + steps: + - name: Check out the repo + uses: actions/checkout@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_PASS }} + + - name: Build r4userembedding and push Docker image + uses: docker/build-push-action@v3 + with: + push: true + tags: beclab/r4userembedding:${{ github.event.inputs.tags }}-amd64 + file: Dockerfile.r4userembedding + platforms: linux/amd64 + + publish_dockerhub_arm64: + runs-on: self-hosted + steps: + - name: Check out the repo + uses: actions/checkout@v3 + + - name: Log in to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_PASS }} + + - name: Build r4userembedding and push Docker image + uses: docker/build-push-action@v3 + with: + push: true + tags: beclab/r4userembedding:${{ github.event.inputs.tags }}-arm64 + file: Dockerfile.r4userembedding + platforms: linux/arm64 + + publish_manifest: + needs: + - publish_dockerhub_amd64 + - publish_dockerhub_arm64 + runs-on: ubuntu-latest + steps: + - name: Log in to Docker Hub + uses: docker/login-action@v2 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_PASS }} + + - name: Push manifest + run: | + docker manifest create beclab/r4userembedding:${{ github.event.inputs.tags }} --amend beclab/r4userembedding:${{ github.event.inputs.tags }}-amd64 --amend beclab/r4userembedding:${{ github.event.inputs.tags }}-arm64 + docker manifest push beclab/r4userembedding:${{ github.event.inputs.tags }} \ No newline at end of file diff --git a/Dockerfile.r4userembedding b/Dockerfile.r4userembedding index 3d6a9a4..687adbc 100644 --- a/Dockerfile.r4userembedding +++ b/Dockerfile.r4userembedding @@ -2,7 +2,7 @@ FROM ubuntu:jammy As builder # RUN apk add openssl-dev musl-dev g++ -RUN mkdir -p /bertv2userembedding && \ +RUN mkdir -p /userembedding && \ apt update && \ apt install curl -y && \ curl https://sh.rustup.rs -sSf | bash -s -- -y && \ @@ -12,9 +12,9 @@ RUN mkdir -p /bertv2userembedding && \ apt-get install pkg-config -y ENV PATH="/root/.cargo/bin:${PATH}" -WORKDIR /bertv2userembedding -COPY user-embedding/src /bertv2userembedding/src -COPY user-embedding/Cargo.toml /bertv2userembedding/Cargo.toml +WORKDIR /userembedding +COPY user-embedding/src /userembedding/src +COPY user-embedding/Cargo.toml /userembedding/Cargo.toml RUN cargo build --release @@ -24,10 +24,10 @@ FROM ubuntu:jammy # Import from builder. -WORKDIR /bertv2userembedding +WORKDIR /userembedding # Copy our build -COPY --from=builder /bertv2userembedding/target/release/bertv2userembedding ./ +COPY --from=builder /userembedding/target/release/ ./ -CMD ["/bertv2userembedding/bertv2userembedding"] \ No newline at end of file +CMD ["/userembedding/userembedding"] \ No newline at end of file diff --git a/user-embedding/Cargo.lock b/user-embedding/Cargo.lock index 656efcf..48d0451 100755 --- a/user-embedding/Cargo.lock +++ b/user-embedding/Cargo.lock @@ -1684,6 +1684,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + [[package]] name = "num-complex" version = "0.4.4" @@ -1694,6 +1704,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-conv" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" + [[package]] name = "num-integer" version = "0.1.45" @@ -1817,6 +1833,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + [[package]] name = "parking_lot" version = "0.12.1" @@ -2450,6 +2472,15 @@ dependencies = [ "digest", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "signal-hook-registry" version = "1.4.1" @@ -2664,6 +2695,16 @@ dependencies = [ "syn 2.0.39", ] +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + [[package]] name = "threadpool" version = "1.8.1" @@ -2675,12 +2716,13 @@ dependencies = [ [[package]] name = "time" -version = "0.3.30" +version = "0.3.36" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4a34ab300f2dee6e562c10a046fc05e358b29f9bf92277f30c3c8d82275f6f5" +checksum = "5dfd88e563464686c916c7e46e623e520ddc6d79fa6641390f2e3fa86e83e885" dependencies = [ "deranged", "itoa", + "num-conv", "powerfmt", "serde", "time-core", @@ -2695,10 +2737,11 @@ checksum = "ef927ca75afb808a4d64dd374f00a2adf8d0fcff8e7b184af886c3c87ec4a3f3" [[package]] name = "time-macros" -version = "0.2.15" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ad70d68dba9e1f8aceda7aa6711965dfec1cac869f311a51bd08b3a2ccbce20" +checksum = "3f252a68540fde3a3877aeea552b832b40ab9a69e318efd078774a01ddee1ccf" dependencies = [ + "num-conv", "time-core", ] @@ -2850,6 +2893,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", ] [[package]] @@ -3036,8 +3105,11 @@ dependencies = [ "serde_json", "text-splitter", "threadpool", + "time", "tokenizers", "tokio", + "tracing", + "tracing-subscriber", ] [[package]] @@ -3056,6 +3128,12 @@ dependencies = [ "serde", ] +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/user-embedding/Cargo.toml b/user-embedding/Cargo.toml index 84cc402..bb12a7c 100755 --- a/user-embedding/Cargo.toml +++ b/user-embedding/Cargo.toml @@ -30,7 +30,10 @@ tokio = {version = "1.33.0",features = ["full"]} ndarray = "0.15.6" ndarray-rand = "0.14.0" reqwest = {version="0.11.23",features = ["json"] } +time = "0.3.36" +tracing = "0.1.40" +tracing-subscriber = "0.3.18" [[bin]] -name = "bertv2userembedding" -path = "src/bertv2_userembedding.rs" +name = "userembedding" +path = "src/userembedding.rs" diff --git a/user-embedding/infra/bertv2_userembedding/Dockerfile b/user-embedding/infra/userembedding/Dockerfile similarity index 58% rename from user-embedding/infra/bertv2_userembedding/Dockerfile rename to user-embedding/infra/userembedding/Dockerfile index 3d6a9a4..687adbc 100644 --- a/user-embedding/infra/bertv2_userembedding/Dockerfile +++ b/user-embedding/infra/userembedding/Dockerfile @@ -2,7 +2,7 @@ FROM ubuntu:jammy As builder # RUN apk add openssl-dev musl-dev g++ -RUN mkdir -p /bertv2userembedding && \ +RUN mkdir -p /userembedding && \ apt update && \ apt install curl -y && \ curl https://sh.rustup.rs -sSf | bash -s -- -y && \ @@ -12,9 +12,9 @@ RUN mkdir -p /bertv2userembedding && \ apt-get install pkg-config -y ENV PATH="/root/.cargo/bin:${PATH}" -WORKDIR /bertv2userembedding -COPY user-embedding/src /bertv2userembedding/src -COPY user-embedding/Cargo.toml /bertv2userembedding/Cargo.toml +WORKDIR /userembedding +COPY user-embedding/src /userembedding/src +COPY user-embedding/Cargo.toml /userembedding/Cargo.toml RUN cargo build --release @@ -24,10 +24,10 @@ FROM ubuntu:jammy # Import from builder. -WORKDIR /bertv2userembedding +WORKDIR /userembedding # Copy our build -COPY --from=builder /bertv2userembedding/target/release/bertv2userembedding ./ +COPY --from=builder /userembedding/target/release/ ./ -CMD ["/bertv2userembedding/bertv2userembedding"] \ No newline at end of file +CMD ["/userembedding/userembedding"] \ No newline at end of file diff --git a/user-embedding/infra/bertv2_userembedding/build.sh b/user-embedding/infra/userembedding/build.sh similarity index 90% rename from user-embedding/infra/bertv2_userembedding/build.sh rename to user-embedding/infra/userembedding/build.sh index 5626f7c..c7c6435 100644 --- a/user-embedding/infra/bertv2_userembedding/build.sh +++ b/user-embedding/infra/userembedding/build.sh @@ -14,4 +14,4 @@ PREFIX=beclab docker build \ -f ${DOCKER_FILE_PATH} \ - -t ${PREFIX}/bertv2userembedding $root_dir \ No newline at end of file + -t ${PREFIX}/userembedding $root_dir \ No newline at end of file diff --git a/user-embedding/src/bertv2.rs b/user-embedding/src/bertcommon.rs similarity index 74% rename from user-embedding/src/bertv2.rs rename to user-embedding/src/bertcommon.rs index a8cb8c2..5577a55 100644 --- a/user-embedding/src/bertv2.rs +++ b/user-embedding/src/bertcommon.rs @@ -12,9 +12,10 @@ use crate::{ knowledge_base_api, }; -const HUGGING_FACE_MODEL_NAME: &str = "sentence-transformers/all-MiniLM-L6-v2"; -const HUGGING_FACE_MODEL_REVISION: &str = "refs/pr/21"; -pub const BERTV2_EMBEDDING_DIMENSION: usize = 384; +// const HUGGING_FACE_MODEL_NAME: &str = "sentence-transformers/all-MiniLM-L6-v2"; +// const HUGGING_FACE_MODEL_REVISION: &str = "refs/pr/21"; +pub const BERT_V2_EMBEDDING_DIMENSION: usize = 384; +pub const BERT_V3_EMBEDDING_DIMENSION: usize = 384; pub async fn calculate_single_entry_pure( current_entry_id: &String, @@ -71,13 +72,24 @@ pub async fn calculate_single_entry_pure( async fn calculate_userembedding() -> AnyhowResult { let current_source_name: String = env::var("TERMINUS_RECOMMEND_SOURCE_NAME") .expect("TERMINUS_RECOMMEND_SOURCE_NAME env not found."); - - let cumulative_embedding_data: [f32; BERTV2_EMBEDDING_DIMENSION] = - [0f32; BERTV2_EMBEDDING_DIMENSION]; - let mut cumulative_tensor = Tensor::new(&cumulative_embedding_data, &Device::Cpu)?; - - let default_model = HUGGING_FACE_MODEL_NAME.to_string(); - let default_revision = HUGGING_FACE_MODEL_REVISION.to_string(); + let embedding_method: String =std::env::var("EMBEDDING_METHOD").expect("EMBEDDING_METHOD not exist"); + let model_related_info: &embedding_common::ModelInfoField = embedding_common::MODEL_RELATED_INFO_MAP.get(embedding_method.as_str()).unwrap(); + let mut option_cumulative_tensor: Option = None; + if model_related_info.model_name == "bert_v2"{ + let cumulative_embedding_data: [f32; BERT_V2_EMBEDDING_DIMENSION] = + [0f32; BERT_V2_EMBEDDING_DIMENSION]; + option_cumulative_tensor = Some(Tensor::new(&cumulative_embedding_data, &Device::Cpu)?); + }else if model_related_info.model_name == "bert_v3"{ + let cumulative_embedding_data: [f32; BERT_V3_EMBEDDING_DIMENSION] = + [0f32; BERT_V3_EMBEDDING_DIMENSION]; + option_cumulative_tensor = Some(Tensor::new(&cumulative_embedding_data, &Device::Cpu)?); + }else{ + tracing::error!("embedding method {} not exist", embedding_method); + return Err(AnyhowError::msg("embedding method not exist")); + } + let mut cumulative_tensor: Tensor = option_cumulative_tensor.unwrap(); + let default_model: String = model_related_info.hugging_face_model_name.to_string(); + let default_revision: String = model_related_info.hugging_face_model_revision.to_string(); let (model, mut tokenizer) = embedding_common::build_model_and_tokenizer(default_model, default_revision).unwrap(); let current_tokenizer: &TokenizerImplSimple = tokenizer @@ -109,7 +121,7 @@ async fn calculate_userembedding() -> AnyhowResult { let current_algorithm_tensor_option = embedding_common::retrieve_current_algorithm_impression_knowledge( current_source_name.clone(), - BERTV2_EMBEDDING_DIMENSION, + model_related_info.embedding_dimension, ) .await; if let Some(current_algorithm_tensor) = current_algorithm_tensor_option { @@ -134,11 +146,13 @@ async fn calculate_userembedding() -> AnyhowResult { } pub async fn execute_bertv2_user_embedding() { + let embedding_method: String =std::env::var("EMBEDDING_METHOD").expect("EMBEDDING_METHOD not exist"); + let model_related_info: &embedding_common::ModelInfoField = embedding_common::MODEL_RELATED_INFO_MAP.get(embedding_method.as_str()).unwrap(); let user_embedding: Tensor = calculate_userembedding() .await .expect("calculate user embedding fail"); let original_user_embedding = - embedding_common::retrieve_user_embedding_through_knowledge(BERTV2_EMBEDDING_DIMENSION) + embedding_common::retrieve_user_embedding_through_knowledge(model_related_info.embedding_dimension) .await .expect("retrieve user embedding through knowledge base fail"); let new_user_embedding_result = user_embedding.add(&original_user_embedding); @@ -167,8 +181,11 @@ mod bertv2test { async fn test_calculate_single_entry() { // cargo test bertv2test::test_calculate_single_entry common_test_operation::init_env(); - let default_model = HUGGING_FACE_MODEL_NAME.to_string(); - let default_revision = HUGGING_FACE_MODEL_REVISION.to_string(); + let embedding_method: String =std::env::var("EMBEDDING_METHOD").expect("EMBEDDING_METHOD not exist"); + let model_related_info: &embedding_common::ModelInfoField = embedding_common::MODEL_RELATED_INFO_MAP.get(embedding_method.as_str()).unwrap(); + + let default_model: String = model_related_info.hugging_face_model_name.to_string(); + let default_revision: String = model_related_info.hugging_face_model_revision.to_string(); let (model, mut tokenizer) = embedding_common::build_model_and_tokenizer(default_model, default_revision).unwrap(); let current_tokenizer: &TokenizerImplSimple = tokenizer @@ -190,7 +207,7 @@ mod bertv2test { #[tokio::test] async fn test_set_init_user_embedding() { common_test_operation::init_env(); - let init_embedding = embedding_common::init_user_embedding(BERTV2_EMBEDDING_DIMENSION); + let init_embedding = embedding_common::init_user_embedding(384); embedding_common::set_user_embedding_knowledgebase(&init_embedding).await; } } diff --git a/user-embedding/src/bertv2_userembedding.rs b/user-embedding/src/bertv2_userembedding.rs deleted file mode 100644 index 504a17c..0000000 --- a/user-embedding/src/bertv2_userembedding.rs +++ /dev/null @@ -1,12 +0,0 @@ -use log::debug as logdebug; -use userembedding::{bertv2, common}; - -#[tokio::main(flavor = "multi_thread", worker_threads = 4)] -async fn main() { - common::init_logger(); - std::env::var("TERMINUS_RECOMMEND_SOURCE_NAME") - .expect("TERMINUS_RECOMMEND_SOURCE_NAME not exist"); - std::env::var("KNOWLEDGE_BASE_API_URL").expect("KNOWLEDGE_BASE_API_URL not exist"); - bertv2::execute_bertv2_user_embedding().await; - logdebug!("execute bertv2_user_embedding compelete"); -} diff --git a/user-embedding/src/common.rs b/user-embedding/src/common.rs index 5e63670..7890933 100755 --- a/user-embedding/src/common.rs +++ b/user-embedding/src/common.rs @@ -8,6 +8,14 @@ use std::{ pub const ENTRY_COLLECTION_NAME: &str = "entries"; pub const IMPRESSIONS_COLLECTION_NAME: &str = "impressions"; + +pub fn init_tracing() { + tracing_subscriber::fmt() + .with_line_number(true) + .with_file(true) + .init(); + tracing::info!("init tracing log successfully"); +} pub fn init_logger() { let env: env_logger::Env<'_> = env_logger::Env::default().filter_or(env_logger::DEFAULT_FILTER_ENV, "debug"); diff --git a/user-embedding/src/common_test_operation.rs b/user-embedding/src/common_test_operation.rs index 34a392d..102d3e1 100644 --- a/user-embedding/src/common_test_operation.rs +++ b/user-embedding/src/common_test_operation.rs @@ -10,7 +10,8 @@ pub fn init_env() { "TERMINUS_RECOMMEND_REDIS_PASSOWRD", "terminusrecommendredis123", ); - env::set_var("KNOWLEDGE_BASE_API_URL", "http://52.202.37.138:3010") + env::set_var("KNOWLEDGE_BASE_API_URL", "http://52.202.37.138:3010"); + env::set_var("EMBEDDING_METHOD", "bert_v2"); // env::set_var("TERMINUS_RECOMMEND_REDIS_URI", "redis://127.0.0.1:6381"); // env::set_var("TERMINUS_RECOMMEND_REDIS_URI", "redis://:terminusrecommendredis123@127.0.0.1:6381"); } diff --git a/user-embedding/src/embedding_common.rs b/user-embedding/src/embedding_common.rs index 6badb5f..dffe864 100755 --- a/user-embedding/src/embedding_common.rs +++ b/user-embedding/src/embedding_common.rs @@ -328,12 +328,48 @@ pub async fn retrieve_current_algorithm_impression_knowledge( return Some(cumulative_tensor); } +pub struct ModelInfoField { + pub model_name: &'static str, + pub hugging_face_model_name: &'static str, + pub hugging_face_model_revision: &'static str, + pub embedding_dimension: usize, + +} + +lazy_static! { + pub static ref MODEL_RELATED_INFO_MAP: HashMap<&'static str, ModelInfoField> = { + let mut m = HashMap::new(); + m.insert( + "bert_v2", + ModelInfoField { + model_name: "bert_v2", + hugging_face_model_name: "sentence-transformers/all-MiniLM-L6-v2", + hugging_face_model_revision: "refs/pr/21", + embedding_dimension: 384, + }, + + ); + m.insert( + "bert_v3", + ModelInfoField { + model_name: "bert_v3", + hugging_face_model_name: "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", + hugging_face_model_revision: "refs/heads/main" , + embedding_dimension: 384, + }, + + ); + m + }; +} + + #[cfg(test)] mod embeddingcommontest { use anyhow::Error as AnyhowError; - use crate::{bertv2, common_test_operation, embedding_common}; + use crate::{bertcommon, common_test_operation, embedding_common}; use super::calculate_one_sentence; use super::*; @@ -346,7 +382,7 @@ mod embeddingcommontest { let current_tensor = retrieve_current_algorithm_impression_knowledge( source_name, - bertv2::BERTV2_EMBEDDING_DIMENSION, + 384, ) .await .expect("add cumulative tensor fail"); diff --git a/user-embedding/src/knowledge_base_api.rs b/user-embedding/src/knowledge_base_api.rs index e887e2e..ba62b42 100644 --- a/user-embedding/src/knowledge_base_api.rs +++ b/user-embedding/src/knowledge_base_api.rs @@ -594,13 +594,13 @@ mod knowledgebaseapitest { use super::*; use crate::common_test_operation; use crate::embedding_common; - use crate::{bertv2, embedding_common::init_user_embedding}; + use crate::{bertcommon, embedding_common::init_user_embedding}; #[tokio::test] async fn test_set_user_embedding_str() { // common::init_logger(); common_test_operation::init_env(); - let init_tensor = init_user_embedding(bertv2::BERTV2_EMBEDDING_DIMENSION); + let init_tensor = init_user_embedding(384); embedding_common::set_user_embedding_knowledgebase(&init_tensor).await; // set_user_embedding_str_through_knowledge(String::from("1.1;1.1;1.2;3"), 3).await; let result = get_user_embedding_str().await; diff --git a/user-embedding/src/lib.rs b/user-embedding/src/lib.rs index 918db1c..8633915 100755 --- a/user-embedding/src/lib.rs +++ b/user-embedding/src/lib.rs @@ -1,6 +1,8 @@ -pub mod bertv2; +pub mod bertcommon; pub mod common; pub mod common_test_operation; pub mod embedding_common; pub mod entity; pub mod knowledge_base_api; +#[macro_use] +extern crate lazy_static; diff --git a/user-embedding/src/userembedding.rs b/user-embedding/src/userembedding.rs new file mode 100644 index 0000000..f3d3b09 --- /dev/null +++ b/user-embedding/src/userembedding.rs @@ -0,0 +1,19 @@ +use log::debug as logdebug; +use userembedding::{bertcommon, common, embedding_common::MODEL_RELATED_INFO_MAP}; + + +#[tokio::main(flavor = "multi_thread", worker_threads = 4)] +async fn main() { + common::init_logger(); + //common::init_tracing(); + std::env::var("TERMINUS_RECOMMEND_SOURCE_NAME") + .expect("TERMINUS_RECOMMEND_SOURCE_NAME not exist"); + let embedding_method: String =std::env::var("EMBEDDING_METHOD").expect("EMBEDDING_METHOD not exist"); + std::env::var("KNOWLEDGE_BASE_API_URL").expect("KNOWLEDGE_BASE_API_URL not exist"); + if MODEL_RELATED_INFO_MAP.contains_key(embedding_method.as_str()) == false { + tracing::error!("embedding method {} not exist", embedding_method); + return; + } + bertcommon::execute_bertv2_user_embedding().await; + logdebug!("execute bertv2_user_embedding compelete"); +}