Skip to content

Commit

Permalink
fix: fix lost bin issue
Browse files Browse the repository at this point in the history
  • Loading branch information
phodal committed Nov 27, 2023
1 parent 20e3121 commit f25f82e
Show file tree
Hide file tree
Showing 2 changed files with 191 additions and 0 deletions.
61 changes: 61 additions & 0 deletions enfer_grpc/src/bin/client.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
use std::io::{Cursor, Read};

use async_stream::stream;
use tokenizer::tokenizer_client::TokenizerClient;
use tokenizer::EncodeRequest;

pub mod tokenizer {
tonic::include_proto!("tokenizer");
}

const model: &[u8] = include_bytes!("../../../model/model.onnx");
const tok: &[u8] = include_bytes!("../../../model/tokenizer.json");

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let mut client = TokenizerClient::connect("http://[::1]:50051").await?;


let mut cursor = Cursor::new(tok);
let mut buf = [0u8; 1024 * 8];

let response = client.set_tokenizer_json(stream! {
while let Ok(n) = cursor.read(&mut buf) {
if n == 0 {
break;
}

yield tokenizer::TokenizerJson {
json: buf[..n].to_vec(),
};
}
}).await?;

println!("tokenizer RESPONSE={:?}", response);

let mut cursor = Cursor::new(model);
let response = client.set_model(stream! {
while let Ok(n) = cursor.read(&mut buf) {
if n == 0 {
break;
}

yield tokenizer::Model {
model: buf[..n].to_vec(),
};
}
}).await?;
println!("model RESPONSE={:?}", response);


let response = client.init_model(()).await ;


let request = tonic::Request::new(EncodeRequest {
text: "Tonic".into(),
});

println!("RESPONSE={:?}", response);

Ok(())
}
130 changes: 130 additions & 0 deletions enfer_grpc/src/bin/server.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
use std::pin::Pin;
use tokio::sync::Mutex;
use clap::Parser;
use tonic::{transport::Server, Request, Response, Status, Streaming};
use tonic::codegen::tokio_stream::StreamExt;

use tokenizer::tokenizer_server::{Tokenizer, TokenizerServer};
use tokenizer::{EncodeRequest, EncodeReply, GeneralResponse, Model, TokenizerJson};

use inference_core::embedding::Semantic;
use inference_grpc::inference_args::InferenceArgs;

pub mod tokenizer {
tonic::include_proto!("tokenizer");
}

#[derive(Default)]
pub struct MyTokenizer {
sema: Mutex<Option<Pin<Box<Semantic>>>>,

tokenzier: Mutex<Vec<u8>>,
model: Mutex<Vec<u8>>,
}

#[tonic::async_trait]
impl Tokenizer for MyTokenizer {
async fn set_tokenizer_json(&self, reqeust: Request<Streaming<TokenizerJson>>) -> Result<Response<GeneralResponse>, Status> {
let mut t = self.tokenzier.lock().await;
t.clear();

let mut stream = reqeust.into_inner();
while let Some(json) = stream.next().await {
let json = match json {
Ok(j) => j,

Err(e) => return Ok(Response::new(GeneralResponse{
success: false,
error: format!("json error: {}", e).into(),
})),
};
t.extend(json.json);
}


Ok(Response::new(GeneralResponse{
success: true,
error: None,
}))
}

async fn set_model(&self, reqeust: Request<Streaming<Model>>) -> Result<Response<GeneralResponse>, Status> {
let mut t = self.model.lock().await;
t.clear();

let mut stream = reqeust.into_inner();
while let Some(model) = stream.next().await {
let model = match model {
Ok(j) => j,

Err(e) => return Ok(Response::new(GeneralResponse{
success: false,
error: format!("model error: {}", e).into(),
})),
};
t.extend(model.model);
}


Ok(Response::new(GeneralResponse{
success: true,
error: None,
}))

}

async fn init_model(&self, _: tonic::Request<()>) -> Result<Response<GeneralResponse>, Status> {
let model = self.model.lock().await;
let tokenizer = self.tokenzier.lock().await;


let sema = match Semantic::initialize(model.clone(), tokenizer.clone()).await {
Ok(t) => t,
Err(e) => return Ok(Response::new(GeneralResponse{
success: false,
error: format!("sma init failed: {}", e).into(),
})),
};

{
let mut s = self.sema.lock().await;
s.replace(sema);
}

Ok(Response::new(GeneralResponse{
success: true,
error: None,
}))

}

async fn encode(
&self,
request: Request<EncodeRequest>,
) -> Result<Response<EncodeReply>, Status> {
println!("Got a request from {:?}", request.remote_addr());

let reply = tokenizer::EncodeReply {
text: format!("Hello {}!", request.into_inner().text),
};

Ok(Response::new(reply))
}
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = InferenceArgs::parse();
let addr = format!("[::1]:{}", args.port.unwrap_or_else(|| "50051".to_string())).parse().unwrap();
let greeter = MyTokenizer::default();

println!("GreeterServer listening on {}", addr);


Server::builder()
.add_service(TokenizerServer::new(greeter))
.serve(addr)
.await?;

Ok(())
}

0 comments on commit f25f82e

Please sign in to comment.