frostyripper1
464 lines
16 KiB
Rust
464 lines
16 KiB
Rust
// src/training/websearch.rs
|
||
#![allow(dead_code)]
|
||
|
||
use crate::data_ls::schema::EnglishSample;
|
||
use anyhow::{anyhow, Result};
|
||
use ollama_rs::generation::chat::request::ChatMessageRequest;
|
||
use ollama_rs::generation::chat::ChatMessage;
|
||
use ollama_rs::generation::completion::request::GenerationRequest;
|
||
use ollama_rs::generation::parameters::{FormatType, KeepAlive};
|
||
use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest};
|
||
use ollama_rs::models::ModelOptions;
|
||
use ollama_rs::Ollama;
|
||
use reqwest::{header, Client};
|
||
use scraper::{Html, Selector};
|
||
use std::fs::OpenOptions;
|
||
use std::io::Write;
|
||
use std::path::Path;
|
||
use urlencoding::encode;
|
||
|
||
fn extract_json_from_response(text: &str) -> Option<String> {
|
||
let trimmed = text.trim();
|
||
let without_fence = trimmed
|
||
.strip_prefix("```json")
|
||
.or_else(|| trimmed.strip_prefix("```json "))
|
||
.or_else(|| trimmed.strip_prefix("```json\n"))
|
||
.or_else(|| trimmed.strip_prefix("```"))
|
||
.unwrap_or(trimmed)
|
||
.trim_start();
|
||
|
||
let without_trailing = without_fence
|
||
.strip_suffix("```")
|
||
.or_else(|| without_fence.strip_suffix("```\n"))
|
||
.or_else(|| without_fence.strip_suffix("``` "))
|
||
.unwrap_or(without_fence)
|
||
.trim();
|
||
|
||
let start = without_trailing.find('{')?;
|
||
let end = without_trailing.rfind('}')?;
|
||
if start >= end {
|
||
return None;
|
||
}
|
||
let json_str = &without_trailing[start..=end];
|
||
if json_str.starts_with('{') && json_str.ends_with('}') {
|
||
Some(json_str.to_string())
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
fn is_json_braces_balanced(text: &str) -> bool {
|
||
let mut balance = 0;
|
||
for c in text.chars() {
|
||
if c == '{' {
|
||
balance += 1;
|
||
} else if c == '}' {
|
||
balance -= 1;
|
||
if balance < 0 {
|
||
return false;
|
||
}
|
||
}
|
||
}
|
||
balance == 0
|
||
}
|
||
|
||
fn repair_json_braces(json: &str) -> String {
|
||
let mut repaired = json.trim_end().to_string();
|
||
let mut balance = 0;
|
||
for c in repaired.chars() {
|
||
if c == '{' {
|
||
balance += 1;
|
||
} else if c == '}' {
|
||
balance -= 1;
|
||
}
|
||
}
|
||
while balance > 0 {
|
||
repaired.push('}');
|
||
balance -= 1;
|
||
}
|
||
repaired
|
||
}
|
||
|
||
fn clean_output(output: &str) -> String {
|
||
let without_url_spaces = output
|
||
.replace("http ://", "http://")
|
||
.replace("https ://", "https://")
|
||
.replace(" : / /", "://");
|
||
|
||
without_url_spaces
|
||
.split_whitespace()
|
||
.collect::<Vec<_>>()
|
||
.join(" ")
|
||
}
|
||
|
||
fn validate_sample(sample: &EnglishSample) -> Result<()> {
|
||
let instruction_lower = sample.instruction.to_lowercase();
|
||
let forbidden_phrases = ["structured learning sample", "valid json format"];
|
||
for forbidden in forbidden_phrases {
|
||
if instruction_lower.contains(forbidden) {
|
||
return Err(anyhow!(
|
||
"Parsed instruction contains forbidden meta-language: '{}'",
|
||
forbidden
|
||
));
|
||
}
|
||
}
|
||
|
||
if sample.output.trim().is_empty() {
|
||
return Err(anyhow!(
|
||
"Parsed sample output is empty. The model must produce a direct answer or summary."
|
||
));
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum SearchEngine {
|
||
Dolphin,
|
||
DuckDuckGo,
|
||
Deepseek,
|
||
Qwen,
|
||
Grok,
|
||
}
|
||
|
||
impl std::fmt::Display for SearchEngine {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
write!(
|
||
f,
|
||
"{}",
|
||
match self {
|
||
SearchEngine::Dolphin => "Dolphin Chat",
|
||
SearchEngine::DuckDuckGo => "DuckDuckGo",
|
||
SearchEngine::Deepseek => "Deepseek",
|
||
SearchEngine::Qwen => "Qwen",
|
||
SearchEngine::Grok => "Grok",
|
||
}
|
||
)
|
||
}
|
||
}
|
||
|
||
impl SearchEngine {
|
||
pub fn search_url(&self, query: &str) -> String {
|
||
match self {
|
||
SearchEngine::Dolphin => format!("https://chat.dphn.ai/?q={}", encode(query)),
|
||
SearchEngine::DuckDuckGo => {
|
||
format!("https://html.duckduckgo.com/html/?q={}", encode(query))
|
||
}
|
||
SearchEngine::Deepseek => format!("https://deepseek.ai/search?q={}", encode(query)),
|
||
SearchEngine::Qwen => format!("https://qwen.ai/search?q={}", encode(query)),
|
||
SearchEngine::Grok => format!("https://grok.com/search?q={}", encode(query)),
|
||
}
|
||
}
|
||
}
|
||
|
||
fn create_browser_client() -> Result<Client, reqwest::Error> {
|
||
let mut headers = header::HeaderMap::new();
|
||
headers.insert(
|
||
header::ACCEPT,
|
||
header::HeaderValue::from_static(
|
||
"text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
|
||
),
|
||
);
|
||
headers.insert(
|
||
header::ACCEPT_LANGUAGE,
|
||
header::HeaderValue::from_static("en-US,en;q=0.9"),
|
||
);
|
||
headers.insert(
|
||
header::CONNECTION,
|
||
header::HeaderValue::from_static("keep-alive"),
|
||
);
|
||
headers.insert(
|
||
header::USER_AGENT,
|
||
header::HeaderValue::from_static("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36"),
|
||
);
|
||
Client::builder().default_headers(headers).build()
|
||
}
|
||
|
||
pub async fn web_search(query: &str, engine: SearchEngine) -> Result<String> {
|
||
let url = engine.search_url(query);
|
||
let client = create_browser_client()?;
|
||
let body = client
|
||
.get(&url)
|
||
.header(header::REFERER, "https://chat.dphn.ai/")
|
||
.send()
|
||
.await?
|
||
.error_for_status()?
|
||
.text()
|
||
.await?;
|
||
let doc = Html::parse_document(&body);
|
||
|
||
let snippets: Vec<String> = match engine {
|
||
SearchEngine::Dolphin => {
|
||
let selector = Selector::parse("p, div, span").unwrap();
|
||
let mut snippets: Vec<String> = doc
|
||
.select(&selector)
|
||
.filter_map(|el| {
|
||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||
if text.len() >= 20 && text.len() <= 300 {
|
||
Some(text)
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.take(5)
|
||
.collect();
|
||
|
||
if snippets.is_empty() {
|
||
let body_text = doc.root_element().text().collect::<Vec<_>>().join(" ");
|
||
snippets = body_text
|
||
.split_terminator(|c| c == '.' || c == '!' || c == '?')
|
||
.map(str::trim)
|
||
.filter(|text| text.len() >= 20 && text.len() <= 300)
|
||
.take(5)
|
||
.map(String::from)
|
||
.collect();
|
||
}
|
||
|
||
snippets
|
||
}
|
||
SearchEngine::DuckDuckGo => {
|
||
let selector = Selector::parse(".result__snippet").unwrap();
|
||
doc.select(&selector)
|
||
.take(5)
|
||
.map(|el| el.inner_html())
|
||
.collect()
|
||
}
|
||
SearchEngine::Deepseek => {
|
||
let selector = Selector::parse(".snippet, .search-result, .result, p").unwrap();
|
||
doc.select(&selector)
|
||
.filter_map(|el| {
|
||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||
if text.len() >= 20 && text.len() <= 300 {
|
||
Some(text)
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.take(5)
|
||
.collect()
|
||
}
|
||
SearchEngine::Qwen => {
|
||
let selector = Selector::parse(".search-result, .result, p").unwrap();
|
||
doc.select(&selector)
|
||
.filter_map(|el| {
|
||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||
if text.len() >= 20 && text.len() <= 300 {
|
||
Some(text)
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.take(5)
|
||
.collect()
|
||
}
|
||
SearchEngine::Grok => {
|
||
let selector = Selector::parse(".search-result, .result, p").unwrap();
|
||
doc.select(&selector)
|
||
.filter_map(|el| {
|
||
let text = el.text().collect::<Vec<_>>().join(" ").trim().to_string();
|
||
if text.len() >= 20 && text.len() <= 300 {
|
||
Some(text)
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.take(5)
|
||
.collect()
|
||
}
|
||
};
|
||
|
||
if snippets.is_empty() {
|
||
return Err(anyhow!(
|
||
"{} returned no usable results from {}",
|
||
engine,
|
||
url
|
||
));
|
||
}
|
||
|
||
Ok(format!("{}\n\nSearch page: {}", snippets.join(" "), url))
|
||
}
|
||
|
||
/// Use LLM to convert search result into a structured EnglishSample
|
||
pub async fn derive_sample_from_search(
|
||
query: &str,
|
||
snippets: &str,
|
||
ollama: &Ollama,
|
||
) -> Result<EnglishSample> {
|
||
let prompt = format!(
|
||
"Convert the following search result into a structured learning sample in JSON format.\n\
|
||
Use this exact schema:\n\
|
||
{{\"type\": \"instruction_following\", \"instruction\": \"{}\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}\n\n\
|
||
Query: {}\nSearch result: {}\n\nOnly output valid JSON, no extra text. Fill in the reasoning and output based on the search result. The instruction should be the query itself or a rephrased version. The type is always \"instruction_following\".",
|
||
query, query, snippets
|
||
);
|
||
let req = ChatMessageRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), vec![ChatMessage::user(prompt)])
|
||
.options(
|
||
ModelOptions::default()
|
||
.num_ctx(8192)
|
||
.num_thread(8)
|
||
.num_predict(1024)
|
||
.temperature(0.85),
|
||
)
|
||
.keep_alive(KeepAlive::Indefinitely);
|
||
let resp = ollama.send_chat_messages(req).await?;
|
||
let json_str = resp.message.content;
|
||
let sample: EnglishSample = serde_json::from_str(&json_str)?;
|
||
Ok(sample)
|
||
}
|
||
|
||
/// Convert a large PDF text into a sequence of structured EnglishSample objects.
|
||
pub async fn derive_samples_from_pdf_text(
|
||
pdf_path: &Path,
|
||
text_chunks: &[String],
|
||
ollama: &Ollama,
|
||
model_name: &str,
|
||
) -> Result<Vec<EnglishSample>> {
|
||
let mut samples = Vec::new();
|
||
let total = text_chunks.len();
|
||
for (index, chunk) in text_chunks.iter().enumerate() {
|
||
let prompt = format!(
|
||
"Extract ONE learning sample from the following text in {}.\n\
|
||
Use this exact schema: {{\"type\": \"instruction_following\", \"instruction\": \"...\", \"input\": \"\", \"reasoning\": \"...\", \"output\": \"...\"}}.\n\
|
||
The instruction must be a specific technical question or task based *only* on the text.\n\
|
||
The output must be a direct answer or summary (no empty strings unless truly absent).\n\
|
||
Do NOT include any meta-dialogue, do NOT describe the format, and do NOT output anything except the JSON object.\n\nText:\n{}",
|
||
pdf_path.file_name().unwrap_or_default().to_string_lossy(),
|
||
chunk
|
||
);
|
||
|
||
let opts = ModelOptions::default()
|
||
.num_ctx(8192)
|
||
.num_thread(8)
|
||
.num_predict(2048)
|
||
.temperature(0.85);
|
||
let request = GenerationRequest::new(model_name.to_string(), prompt.clone())
|
||
.format(FormatType::Json)
|
||
.options(opts.clone());
|
||
|
||
let resp = match ollama.generate(request).await {
|
||
Ok(r) => r,
|
||
Err(e) if model_name == "deepseek-r1:7b" => {
|
||
eprintln!(
|
||
"⚠️ deepseek-r1:7b failed, falling back to qwen2.5-coder:1.5b. Error: {}",
|
||
e
|
||
);
|
||
let fallback_req = GenerationRequest::new("qwen2.5-coder:1.5b".to_string(), prompt.clone())
|
||
.format(FormatType::Json)
|
||
.options(opts.clone());
|
||
match ollama.generate(fallback_req).await {
|
||
Ok(r) => r,
|
||
Err(_) => {
|
||
eprintln!(
|
||
"⚠️ qwen2.5-coder:1.5b also failed; falling back to sadiq-bd/llama3.2-3b-uncensored."
|
||
);
|
||
let fallback_req = GenerationRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), prompt)
|
||
.format(FormatType::Json)
|
||
.options(opts);
|
||
ollama.generate(fallback_req).await?
|
||
}
|
||
}
|
||
}
|
||
Err(e) if model_name == "qwen2.5-coder:1.5b" => {
|
||
eprintln!(
|
||
"⚠️ qwen2.5-coder:1.5b failed, falling back to sadiq-bd/llama3.2-3b-uncensored. Error: {}",
|
||
e
|
||
);
|
||
let fallback_req = GenerationRequest::new("sadiq-bd/llama3.2-3b-uncensored".to_string(), prompt)
|
||
.format(FormatType::Json)
|
||
.options(opts);
|
||
ollama.generate(fallback_req).await?
|
||
}
|
||
Err(e) => return Err(e.into()),
|
||
};
|
||
|
||
let raw_response = resp.response;
|
||
let mut json_candidate = extract_json_from_response(&raw_response)
|
||
.unwrap_or_else(|| raw_response.trim().to_string());
|
||
if !is_json_braces_balanced(&json_candidate) {
|
||
json_candidate = repair_json_braces(&json_candidate);
|
||
}
|
||
|
||
let sample: EnglishSample = match serde_json::from_str::<EnglishSample>(&json_candidate) {
|
||
Ok(mut s) => {
|
||
s.output = clean_output(&s.output);
|
||
s
|
||
}
|
||
Err(err) => {
|
||
let repaired_candidate = repair_json_braces(&json_candidate);
|
||
if repaired_candidate != json_candidate {
|
||
serde_json::from_str(&repaired_candidate).map_err(|second_err| {
|
||
anyhow::anyhow!(
|
||
"Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCandidate: {}\nRepaired: {}\nRaw response: {}",
|
||
index + 1,
|
||
total,
|
||
err,
|
||
second_err,
|
||
json_candidate,
|
||
repaired_candidate,
|
||
raw_response
|
||
)
|
||
})?
|
||
} else if let Some(cleaned) = extract_json_from_response(&raw_response) {
|
||
serde_json::from_str(&cleaned).map_err(|second_err| {
|
||
anyhow::anyhow!(
|
||
"Failed to parse JSON from PDF chunk {} of {}: {}\nRetry parse error: {}\nCleaned response: {}\nRaw response: {}",
|
||
index + 1,
|
||
total,
|
||
err,
|
||
second_err,
|
||
cleaned,
|
||
raw_response
|
||
)
|
||
})?
|
||
} else {
|
||
return Err(anyhow::anyhow!(
|
||
"Failed to parse JSON from PDF chunk {} of {}.\nRaw response:\n{}",
|
||
index + 1,
|
||
total,
|
||
raw_response
|
||
));
|
||
}
|
||
}
|
||
};
|
||
|
||
validate_sample(&sample).map_err(|validation_err| {
|
||
anyhow::anyhow!(
|
||
"Validation failed for PDF chunk {} of {}: {}\nCandidate: {}\nRaw response: {}",
|
||
index + 1,
|
||
total,
|
||
validation_err,
|
||
json_candidate,
|
||
raw_response
|
||
)
|
||
})?;
|
||
|
||
samples.push(sample);
|
||
}
|
||
Ok(samples)
|
||
}
|
||
|
||
/// Append a new sample to the dynamic knowledge file
|
||
pub fn append_sample_to_jsonl(sample: &EnglishSample, path: &str) -> Result<()> {
|
||
let mut file = OpenOptions::new().create(true).append(true).open(path)?;
|
||
let line = serde_json::to_string(sample)?;
|
||
writeln!(file, "{}", line)?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Compute embedding for a sample (so we can add it to in‑memory vectors)
|
||
pub async fn compute_sample_embedding(
|
||
sample: &EnglishSample,
|
||
embed_model: &str,
|
||
ollama: &Ollama,
|
||
) -> Result<Vec<f32>> {
|
||
let text = format!(
|
||
"Instruction: {}\nOutput: {}",
|
||
sample.instruction, sample.output
|
||
);
|
||
let req =
|
||
GenerateEmbeddingsRequest::new(embed_model.to_string(), EmbeddingsInput::Single(text))
|
||
.keep_alive(KeepAlive::Indefinitely);
|
||
let resp = ollama.generate_embeddings(req).await?;
|
||
Ok(resp.embeddings.into_iter().next().unwrap_or_default())
|
||
}
|