Skip to content

Commit

Permalink
improve rag
Browse files Browse the repository at this point in the history
  • Loading branch information
santiagomed committed Dec 13, 2023
1 parent 4c9b433 commit 785c566
Show file tree
Hide file tree
Showing 4 changed files with 14 additions and 2 deletions.
8 changes: 7 additions & 1 deletion orca-core/src/llm/quantized.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,11 +122,17 @@ impl Quantized {
self.sample_len = sample_len;
self
}

pub fn with_model(mut self, model: Model) -> Self {
self.which = model;
self
}

pub fn with_seed(mut self, seed: u64) -> Self {
self.seed = seed;
self
}

fn tokenizer(&self) -> anyhow::Result<Tokenizer> {
let tokenizer_path = match &self.tokenizer {
Some(config) => std::path::PathBuf::from(config),
Expand Down Expand Up @@ -314,7 +320,7 @@ impl LLM for Quantized {
Quantized::format_chat_prompt(prompt.to_chat()?)
};

log::info!("prompt:\n{}", &prompt);
log::debug!("prompt:\n{}", &prompt);
let mut result = String::new();
let tokens = tokenizer.encode(prompt, true).map_err(anyhow::Error::msg)?;
if log::log_enabled!(log::Level::Debug) {
Expand Down
4 changes: 4 additions & 0 deletions orca-core/src/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,10 @@ impl ChatBuffer {
memory: ChatPrompt(Vec::new()),
}
}

pub fn from_chat(chat: &ChatPrompt) -> Self {
Self { memory: chat.clone() }
}
}

impl Memory for ChatBuffer {
Expand Down
1 change: 1 addition & 0 deletions orca-core/src/pipeline/simple.rs
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,7 @@ impl<M: LLM + Clone + 'static> Pipeline for LLMPipeline<M> {
let mut locked_memory = memory.lock().await; // Lock the memory
let mem = locked_memory.memory();
mem.save(prompt);
log::debug!("Memory: {}", mem);
self.llm.generate(mem.clone_prompt()).await?
} else {
self.llm.generate(prompt.clone_prompt()).await?
Expand Down
3 changes: 2 additions & 1 deletion orca-core/src/prompt/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -235,14 +235,15 @@ pub trait Prompt: Sync + Send + Display {
/// ```
fn save(&mut self, _data: Box<dyn Prompt>) {
unimplemented!("save not implemented for this prompt type");
// Err(anyhow::anyhow!("save not implemented for this prompt type"))
}

/// Convert the current prompt to a `ChatPrompt`.
///
/// # Returns
/// * `Result<ChatPrompt>` - The `ChatPrompt` representation of the prompt or an error.
fn to_chat(&self) -> Result<ChatPrompt> {
unimplemented!("Unable to convert prompt to ChatPrompt");
Err(anyhow::anyhow!("Unable to convert prompt to ChatPrompt"))
}

/// Clone the current prompt into a Boxed trait object.
Expand Down

0 comments on commit 785c566

Please sign in to comment.