Oooh its better i think I dont know

This commit is contained in:
frostyripper1
2026-05-20 09:22:26 +02:00
parent fe97362aad
commit ee34649e02
11 changed files with 220 additions and 171 deletions
+26 -9
View File
@@ -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
View File
@@ -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);
}