izzilis/src/model.rs

80 lines
2.3 KiB
Rust
Raw Normal View History

use async_std::{io, process::Command};
use futures::{future::BoxFuture, stream::BoxStream, Future, Stream, StreamExt, TryStreamExt};
pub trait SampleModel {
type Error;
type Sample: Future<Output = Result<String, Self::Error>>;
fn get_sample(&self) -> Self::Sample;
}
pub struct GPTSampleModel {
python_command: String,
command_working_path: String,
command_args: Vec<String>,
}
impl SampleModel for GPTSampleModel {
type Error = io::Error;
type Sample = BoxFuture<'static, Result<String, Self::Error>>;
fn get_sample(&self) -> Self::Sample {
let cmd = Command::new(&self.python_command)
.current_dir(&self.command_working_path)
.args(&self.command_args)
.output();
Box::pin(async { Ok(String::from_utf8_lossy(&cmd.await?.stdout).to_string()) })
}
}
impl GPTSampleModel {
pub fn new(
python_command: String,
command_working_path: String,
command_args: Vec<String>,
) -> GPTSampleModel {
Self {
python_command: python_command,
command_working_path: command_working_path,
command_args: command_args,
}
}
}
pub trait SampleModelExt: SampleModel {
type Stream: Stream<Item = Result<String, Self::Error>>;
fn into_stream(self) -> Self::Stream;
}
const SAMPLE_SPLIT_WORD: &str = "<|endoftext|>";
const SAMPLE_SAMPLE_LINE: &str =
"======================================== SAMPLE 1 ========================================";
impl<T: SampleModel + Send + Sync + 'static> SampleModelExt for T
where
Self::Sample: Send,
{
type Stream = BoxStream<'static, Result<String, Self::Error>>;
fn into_stream(self) -> Self::Stream {
Box::pin(
futures::stream::try_unfold(self, |this| async {
Ok(Some((this.get_sample().await?, this)))
})
.map_ok(|samples| {
futures::stream::iter(
samples
.replace(SAMPLE_SAMPLE_LINE, "")
.split(SAMPLE_SPLIT_WORD)
.map(|elem| elem.to_owned())
.collect::<Vec<String>>()
.into_iter()
.map(|elem| Ok(elem.trim().to_owned())),
)
})
.try_flatten(),
)
}
}