Initial commit - AiRust cyber operator
This commit is contained in:
frostyripper1
@@ -0,0 +1,257 @@
|
||||
use crate::training::websearch::derive_samples_from_pdf_text;
|
||||
use anyhow::{anyhow, Context, Result};
|
||||
use lopdf::Document;
|
||||
use ollama_rs::Ollama;
|
||||
use pdf_extract::extract_text;
|
||||
use serde_json;
|
||||
use std::fs::{self, File};
|
||||
use std::io::{self, BufWriter, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Scan the given directory and return all PDF file paths.
|
||||
fn find_pdf_files(input_dir: &Path) -> Result<Vec<PathBuf>> {
|
||||
if !input_dir.is_dir() {
|
||||
return Err(anyhow!("{} is not a directory", input_dir.display()));
|
||||
}
|
||||
|
||||
let mut pdf_files = Vec::new();
|
||||
for entry in fs::read_dir(input_dir)
|
||||
.with_context(|| format!("Failed to read directory {}", input_dir.display()))?
|
||||
{
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if path.is_file() {
|
||||
if let Some(ext) = path.extension() {
|
||||
if ext.to_string_lossy().eq_ignore_ascii_case("pdf") {
|
||||
pdf_files.push(path);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(pdf_files)
|
||||
}
|
||||
|
||||
const MAX_CHUNK_CHARS: usize = 10000;
|
||||
|
||||
/// Split large PDF text into manageable chunks while preserving whole paragraphs.
|
||||
fn split_text_into_chunks(text: &str, max_chars: usize) -> Vec<String> {
|
||||
let mut chunks = Vec::new();
|
||||
let mut current = String::new();
|
||||
|
||||
for paragraph in text.split('\n').map(str::trim).filter(|p| !p.is_empty()) {
|
||||
if current.len() + paragraph.len() + 1 > max_chars {
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
current = String::new();
|
||||
}
|
||||
}
|
||||
|
||||
if paragraph.len() > max_chars {
|
||||
for piece in paragraph.chars().collect::<Vec<_>>().chunks(max_chars) {
|
||||
let piece_str: String = piece.iter().collect();
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
current = String::new();
|
||||
}
|
||||
chunks.push(piece_str);
|
||||
}
|
||||
} else {
|
||||
if !current.is_empty() {
|
||||
current.push(' ');
|
||||
}
|
||||
current.push_str(paragraph);
|
||||
}
|
||||
}
|
||||
|
||||
if !current.is_empty() {
|
||||
chunks.push(current);
|
||||
}
|
||||
|
||||
chunks
|
||||
}
|
||||
|
||||
pub async fn import_pdf_directory(
|
||||
input_dir: &Path,
|
||||
output_file: &Path,
|
||||
ollama: &Ollama,
|
||||
model_name: &str,
|
||||
prompt_each_chunk: bool,
|
||||
) -> Result<bool> {
|
||||
let pdf_files = find_pdf_files(input_dir)?;
|
||||
if pdf_files.is_empty() {
|
||||
return Err(anyhow!("No PDF files found in {}", input_dir.display()));
|
||||
}
|
||||
|
||||
if let Some(parent) = output_file.parent() {
|
||||
fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
let mut writer = BufWriter::new(File::create(output_file)?);
|
||||
println!(
|
||||
"📦 Found {} PDF file(s) in {}",
|
||||
pdf_files.len(),
|
||||
input_dir.display()
|
||||
);
|
||||
|
||||
let mut total_samples = 0;
|
||||
for pdf_path in &pdf_files {
|
||||
println!("📄 Converting {}", pdf_path.display());
|
||||
let text = extract_text(pdf_path)
|
||||
.with_context(|| format!("Failed to extract text from {}", pdf_path.display()))?;
|
||||
|
||||
let pages_count = count_pdf_pages(pdf_path).unwrap_or(0);
|
||||
let chunks = split_text_into_chunks(&text, MAX_CHUNK_CHARS);
|
||||
if pages_count > 0 {
|
||||
println!(
|
||||
" ➜ PDF produced {} chunk(s) from {} pages",
|
||||
chunks.len(),
|
||||
pages_count
|
||||
);
|
||||
} else {
|
||||
println!(
|
||||
" ➜ PDF produced {} chunk(s); could not detect page count",
|
||||
chunks.len()
|
||||
);
|
||||
}
|
||||
|
||||
let (count, completed) =
|
||||
process_pdf_chunks(pdf_path, &chunks, &mut writer, ollama, model_name, pages_count, prompt_each_chunk).await?;
|
||||
total_samples += count;
|
||||
if !completed {
|
||||
println!(
|
||||
"⚠️ Stopping directory conversion after partial output for {}",
|
||||
pdf_path.display()
|
||||
);
|
||||
return Ok(false);
|
||||
}
|
||||
}
|
||||
|
||||
writer.flush()?;
|
||||
println!(
|
||||
"✅ Wrote {} JSONL sample(s) to {}",
|
||||
total_samples,
|
||||
output_file.display()
|
||||
);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub async fn import_pdf_path(
|
||||
input_path: &Path,
|
||||
output_file: &Path,
|
||||
ollama: &Ollama,
|
||||
model_name: &str,
|
||||
prompt_each_chunk: bool,
|
||||
) -> Result<bool> {
|
||||
if input_path.is_dir() {
|
||||
import_pdf_directory(input_path, output_file, ollama, model_name, prompt_each_chunk).await
|
||||
} else if input_path.is_file() {
|
||||
if input_path
|
||||
.extension()
|
||||
.map(|ext| ext.to_string_lossy().eq_ignore_ascii_case("pdf"))
|
||||
.unwrap_or(false)
|
||||
{
|
||||
if let Some(parent) = output_file.parent() {
|
||||
fs::create_dir_all(parent)?;
|
||||
}
|
||||
|
||||
let mut writer = BufWriter::new(File::create(output_file)?);
|
||||
println!("📄 Converting {}", input_path.display());
|
||||
let text = extract_text(input_path)
|
||||
.with_context(|| format!("Failed to extract text from {}", input_path.display()))?;
|
||||
|
||||
let pages_count = count_pdf_pages(input_path).unwrap_or(0);
|
||||
let chunks = split_text_into_chunks(&text, MAX_CHUNK_CHARS);
|
||||
if pages_count > 0 {
|
||||
println!(
|
||||
" ➜ PDF produced {} chunk(s) from {} pages",
|
||||
chunks.len(),
|
||||
pages_count
|
||||
);
|
||||
} else {
|
||||
println!(
|
||||
" ➜ PDF produced {} chunk(s); could not detect page count",
|
||||
chunks.len()
|
||||
);
|
||||
}
|
||||
|
||||
let (total_samples, completed) =
|
||||
process_pdf_chunks(input_path, &chunks, &mut writer, ollama, model_name, pages_count, prompt_each_chunk).await?;
|
||||
writer.flush()?;
|
||||
println!(
|
||||
"✅ Wrote {} JSONL sample(s) to {}",
|
||||
total_samples,
|
||||
output_file.display()
|
||||
);
|
||||
if completed {
|
||||
println!("✅ Done! Full PDF distillation complete.");
|
||||
} else {
|
||||
println!("⚠️ Partial output created; import stopped early.");
|
||||
}
|
||||
Ok(completed)
|
||||
} else {
|
||||
Err(anyhow!(
|
||||
"The selected file is not a PDF: {}",
|
||||
input_path.display()
|
||||
))
|
||||
}
|
||||
} else {
|
||||
Err(anyhow!(
|
||||
"The selected path is neither a PDF file nor a directory: {}",
|
||||
input_path.display()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
async fn process_pdf_chunks(
|
||||
pdf_path: &Path,
|
||||
chunks: &[String],
|
||||
writer: &mut BufWriter<File>,
|
||||
ollama: &Ollama,
|
||||
model_name: &str,
|
||||
pages_count: usize,
|
||||
prompt_each_chunk: bool,
|
||||
) -> Result<(usize, bool)> {
|
||||
let mut total_samples = 0;
|
||||
for (index, chunk) in chunks.iter().enumerate() {
|
||||
let progress = format!("chunk {}/{}", index + 1, chunks.len());
|
||||
let page_label = if pages_count > 0 {
|
||||
let current_page = ((index * pages_count) / chunks.len()) + 1;
|
||||
format!("page {}/{}", current_page, pages_count)
|
||||
} else {
|
||||
"page unknown".to_string()
|
||||
};
|
||||
println!(" ➜ Processing {} ({})", progress, page_label);
|
||||
let samples =
|
||||
derive_samples_from_pdf_text(pdf_path, &[chunk.clone()], ollama, model_name).await?;
|
||||
for sample in samples {
|
||||
let line = serde_json::to_string(&sample)?;
|
||||
writeln!(writer, "{}", line)?;
|
||||
total_samples += 1;
|
||||
}
|
||||
|
||||
if index + 1 < chunks.len() && prompt_each_chunk {
|
||||
if !prompt_continue()? {
|
||||
println!("⚠️ Stopped early. Partial output written to disk.");
|
||||
return Ok((total_samples, false));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok((total_samples, true))
|
||||
}
|
||||
|
||||
fn count_pdf_pages(pdf_path: &Path) -> Result<usize> {
|
||||
let doc = Document::load(pdf_path)?;
|
||||
Ok(doc.get_pages().len())
|
||||
}
|
||||
|
||||
fn prompt_continue() -> Result<bool> {
|
||||
print!("Press Enter to continue to the next chunk, or type 'quit' to stop: ");
|
||||
let stdout = io::stdout();
|
||||
let mut handle = stdout.lock();
|
||||
handle.flush()?;
|
||||
|
||||
let mut input = String::new();
|
||||
io::stdin().read_line(&mut input)?;
|
||||
Ok(!input.trim().eq_ignore_ascii_case("quit"))
|
||||
}
|
||||
Reference in New Issue
Block a user