Oooh its better i think I dont know
This commit is contained in:
+26
-9
@@ -6,7 +6,7 @@ use ollama_rs::Ollama;
|
||||
use rand::seq::IndexedRandom;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::fs;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use std::io::{self, Write};
|
||||
@@ -80,14 +80,18 @@ pub struct GraphEdge {
|
||||
pub relation: String,
|
||||
}
|
||||
|
||||
const EMBED_CACHE_MAX: usize = 10_000;
|
||||
|
||||
pub struct EmbeddingCache {
|
||||
cache: HashMap<u64, Vec<f32>>,
|
||||
order: VecDeque<u64>,
|
||||
}
|
||||
|
||||
impl EmbeddingCache {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cache: HashMap::new(),
|
||||
order: VecDeque::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -96,7 +100,17 @@ impl EmbeddingCache {
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, text: &str, emb: Vec<f32>) {
|
||||
self.cache.insert(hash_text(text), emb);
|
||||
let key = hash_text(text);
|
||||
if self.cache.contains_key(&key) {
|
||||
return;
|
||||
}
|
||||
if self.cache.len() >= EMBED_CACHE_MAX {
|
||||
if let Some(old) = self.order.pop_front() {
|
||||
self.cache.remove(&old);
|
||||
}
|
||||
}
|
||||
self.cache.insert(key, emb);
|
||||
self.order.push_back(key);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -343,7 +357,9 @@ impl KnowledgeIndex {
|
||||
let _ = fs::create_dir_all(parent);
|
||||
}
|
||||
if let Ok(data) = serde_json::to_string(&existing) {
|
||||
let _ = fs::write(ARCHIVE_PATH, data);
|
||||
if let Err(e) = fs::write(ARCHIVE_PATH, &data) {
|
||||
eprintln!("archive write failed: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
self.archived = Some(existing);
|
||||
@@ -484,7 +500,7 @@ impl KnowledgeIndex {
|
||||
(idx, hybrid)
|
||||
})
|
||||
.collect();
|
||||
ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
ranked.into_iter().map(|(idx, _)| idx).collect()
|
||||
}
|
||||
|
||||
@@ -549,7 +565,7 @@ impl KnowledgeIndex {
|
||||
.enumerate()
|
||||
.map(|(idx, emb)| (idx, cosine_similarity(query_emb, emb)))
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
scored
|
||||
.into_iter()
|
||||
.filter(|(_, s)| *s >= threshold)
|
||||
@@ -578,7 +594,7 @@ impl KnowledgeIndex {
|
||||
.enumerate()
|
||||
.map(|(idx, emb)| (idx, cosine_similarity(query_emb, emb)))
|
||||
.collect();
|
||||
s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
s
|
||||
} else {
|
||||
let mut centroid_scores: Vec<(usize, f32)> = self
|
||||
@@ -587,7 +603,8 @@ impl KnowledgeIndex {
|
||||
.enumerate()
|
||||
.map(|(i, c)| (i, cosine_similarity(query_emb, c)))
|
||||
.collect();
|
||||
centroid_scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
centroid_scores
|
||||
.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let nprobe = (self.nlist / 5).max(1);
|
||||
let top_centroids: Vec<usize> = centroid_scores
|
||||
.into_iter()
|
||||
@@ -601,7 +618,7 @@ impl KnowledgeIndex {
|
||||
.filter(|(idx, _)| top_centroids.contains(&self.assignments[*idx]))
|
||||
.map(|(idx, emb)| (idx, cosine_similarity(query_emb, emb)))
|
||||
.collect();
|
||||
s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
s
|
||||
};
|
||||
|
||||
@@ -688,7 +705,7 @@ fn nearest_centroid(vec: &[f32], centroids: &[Vec<f32>]) -> usize {
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| (i, cosine_similarity(vec, c)))
|
||||
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
|
||||
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal))
|
||||
.map(|(i, _)| i)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
+29
-8
@@ -308,6 +308,10 @@ async fn handle_chat(
|
||||
let mut h = hist_arc.lock().await;
|
||||
h.push(ChatMessage::user(query_for_hist));
|
||||
h.push(ChatMessage::assistant(full_reply));
|
||||
if h.len() > 100 {
|
||||
let excess = h.len() - 100;
|
||||
h.drain(..excess);
|
||||
}
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
let _ = tx
|
||||
@@ -385,7 +389,9 @@ async fn maybe_save_cache(knowledge: &KnowledgeIndex) {
|
||||
if now - last < 5 {
|
||||
return;
|
||||
}
|
||||
let _ = knowledge.save_cache();
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
}
|
||||
LAST_CACHE_SAVE.store(now, std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
|
||||
@@ -552,7 +558,9 @@ async fn handle_learn_search(
|
||||
}
|
||||
};
|
||||
|
||||
let _ = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl");
|
||||
if let Err(e) = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl") {
|
||||
eprintln!("append_sample failed: {e}");
|
||||
}
|
||||
|
||||
let sample_text = format!(
|
||||
"Instruction: {}\nOutput: {}",
|
||||
@@ -637,7 +645,10 @@ async fn handle_learn_auto(
|
||||
}
|
||||
};
|
||||
|
||||
let _ = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl");
|
||||
if let Err(e) = append_sample_to_jsonl(&sample, "output/random_new_knowledge.jsonl")
|
||||
{
|
||||
eprintln!("append_sample failed: {e}");
|
||||
}
|
||||
|
||||
let sample_text = format!(
|
||||
"Instruction: {}\nOutput: {}",
|
||||
@@ -683,7 +694,9 @@ async fn handle_learn_auto(
|
||||
|
||||
// Save cache once after all parallel work
|
||||
let knowledge = state.knowledge.lock().await;
|
||||
let _ = knowledge.save_cache();
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
}
|
||||
drop(knowledge);
|
||||
|
||||
let learned = results.iter().filter(|r| r.get("error").is_none()).count();
|
||||
@@ -894,7 +907,9 @@ async fn handle_pdf_import(
|
||||
let existing = knowledge.len();
|
||||
knowledge.extend_from_samples(clean, &source_name);
|
||||
if knowledge.len() > existing {
|
||||
let _ = knowledge.save_cache();
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
}
|
||||
}
|
||||
}
|
||||
drop(knowledge);
|
||||
@@ -1082,7 +1097,9 @@ async fn handle_epub_import(
|
||||
for s in &samples {
|
||||
let line = serde_json::to_string(s).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
writeln!(out_file, "{}", line).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
let _ = append_sample_to_jsonl(s, "output/random_new_knowledge.jsonl");
|
||||
if let Err(e) = append_sample_to_jsonl(s, "output/random_new_knowledge.jsonl") {
|
||||
eprintln!("append_sample failed: {e}");
|
||||
}
|
||||
count += 1;
|
||||
}
|
||||
|
||||
@@ -1091,7 +1108,9 @@ async fn handle_epub_import(
|
||||
crate::pipeline::preprocess::remove_empty(samples),
|
||||
&format!("epub:{}", filename),
|
||||
);
|
||||
let _ = knowledge.save_cache();
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
}
|
||||
drop(knowledge);
|
||||
|
||||
let _ = std::fs::remove_file(&epub_path);
|
||||
@@ -1280,7 +1299,9 @@ async fn handle_dataset_generate(
|
||||
}
|
||||
}
|
||||
}
|
||||
let _ = knowledge.save_cache();
|
||||
if let Err(e) = knowledge.save_cache() {
|
||||
eprintln!("save_cache failed: {e}");
|
||||
}
|
||||
drop(knowledge);
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user