Add rag-ingestion module: PDF, Word and HTML readers and chunking strategies compared by retrieval hit rate; metadata in the embedding, content-hash idempotent ingestion, re-embedding cost of an edit
Co-Authored-By: Claude Sonnet 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_0146opDKBYtADXvkc39zdiyw
This commit is contained in:
Claude
@@ -0,0 +1,137 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* The one document every experiment in this module reads: a made-up operations runbook with 12 sections
|
||||
* and 36 paragraphs, plus 37 questions whose answers sit in exactly one paragraph each.
|
||||
*
|
||||
* <p>It is written out as a PDF, a Word file and an HTML page (see the test fixtures), so the readers and
|
||||
* the chunkers are compared on identical content. A question counts as answered when a retrieved chunk
|
||||
* contains its {@code answer} fragment, and {@code CorpusTest} proves every fragment occurs exactly once.
|
||||
*/
|
||||
public final class Corpus {
|
||||
|
||||
public static final String TITLE = "Helios Cloud Platform Operations Runbook";
|
||||
|
||||
public record Section(String title, List<String> paragraphs) {
|
||||
}
|
||||
|
||||
public record Question(String text, String answer) {
|
||||
}
|
||||
|
||||
public static final List<Section> SECTIONS = List.of(
|
||||
new Section("1. Backups and Restore", List.of(
|
||||
"Every production database is snapshotted once a day at 01:30 UTC. Daily snapshots are kept for 35 days, after which they are deleted automatically. Weekly snapshots, taken on Sundays, are kept for 12 weeks. All snapshots are encrypted with AES-256 using keys held in the regional key management service, and each snapshot is copied to the eu-west-2 region within one hour of completion.",
|
||||
"A backup that has never been restored is only a hope. On the first Monday of every quarter the storage team restores the newest snapshot of three randomly chosen databases into an isolated account and runs the integrity checklist against it. A restore test passes only if the restored database starts, replays its write-ahead log without errors and returns the expected row counts for the canary tables.",
|
||||
"Point-in-time recovery is available for the last 7 days. To request a recovery, the service owner opens a ticket in the STORAGE queue with the target timestamp in UTC. The storage team aims to start the recovery within 2 hours of a request and will always restore into a new instance, never over the live one, so that the owner can compare the data before switching traffic.")),
|
||||
new Section("2. Incident Severity and Response", List.of(
|
||||
"Incidents are classified by customer impact. A Sev1 is a full outage of a customer-facing service or any confirmed loss of customer data. A Sev2 is a partial outage or a degradation that affects more than 10 percent of requests. A Sev3 is anything else that needs a human but does not hurt customers right now. The incident commander decides the severity and may raise or lower it at any time during the incident.",
|
||||
"The response targets are fixed: for a Sev1 a responder must be working within 10 minutes of the page, for a Sev2 within 30 minutes and for a Sev3 by the end of the next business day. Sev1 incidents get a status page update within 20 minutes and then every 30 minutes until resolved. The communications lead, not the incident commander, writes those updates.",
|
||||
"Every Sev1 and Sev2 gets a written post-incident review within 5 business days. The review is blameless, lists the contributing factors, and ends with action items that each have an owner and a due date. Action items that are still open after 60 days are escalated to the director of engineering.")),
|
||||
new Section("3. On-Call Rotation", List.of(
|
||||
"Each service has a primary and a secondary on-call engineer. The rotation changes every Thursday at 10:00 UTC, and the outgoing engineer gives a 30 minute verbal handover covering open alerts, risky changes and anything that felt unusual during the week. Nobody may be on call for more than 2 consecutive weeks without written approval from their manager.",
|
||||
"If the primary does not acknowledge a page within 5 minutes it escalates to the secondary, and after a further 5 minutes to the engineering manager. On-call engineers are paid an allowance of 120 euros for each week on the rotation, plus time off in lieu for any night call that lasts longer than 1 hour.",
|
||||
"New engineers shadow the rotation for two full weeks before taking the pager. During shadowing they receive every page but are not expected to act, and they write up what they would have done so that the primary can review it afterwards.")),
|
||||
new Section("4. Deployments and Releases", List.of(
|
||||
"Production deployments are blocked every Friday after 14:00 UTC and for the whole period from 20 December to 3 January. An emergency change during a freeze needs approval from the on-call incident commander and is recorded in the change log with the reason for the exception.",
|
||||
"Every release starts as a canary that receives 5 percent of traffic for 20 minutes. If the error rate of the canary rises more than 2 percent above the baseline in that window the pipeline rolls back automatically without waiting for a human. If the canary stays healthy it is promoted in steps of 25 percent, with a 10 minute pause between steps.",
|
||||
"Database migrations are shipped separately from application code and must be backwards compatible with the previous release, because a rollback restores only the code. A migration that rewrites more than 1 million rows must run as a background job with a throttle and must be announced in the platform channel one day in advance.")),
|
||||
new Section("5. Access Control", List.of(
|
||||
"Engineers reach production only through short-lived certificates issued by the access broker. A certificate is valid for 8 hours and cannot be renewed without signing in again with a hardware security key. Passwords and SMS codes are not accepted for any production system.",
|
||||
"Access rights are reviewed every quarter by the owning team lead, and rights that were not used in the previous 90 days are removed automatically. When an engineer changes teams, their previous team's access is removed on the same day and not at the next review.",
|
||||
"The break-glass account exists for the case where the access broker itself is down. Its token expires after 60 minutes, every use pages the security team immediately, and the engineer who used it must file a written justification within 24 hours.")),
|
||||
new Section("6. Capacity and Autoscaling", List.of(
|
||||
"Services scale on CPU utilisation with a target of 65 percent. After a scale-out the autoscaler waits 180 seconds before it considers another one, and after a scale-in it waits 600 seconds, so that a short burst cannot make the fleet flap up and down.",
|
||||
"Each region keeps reserved headroom of 30 percent above the busiest hour of the last month. Capacity is validated by a load test on the first Wednesday of the month at 3 times the observed peak, and a service that cannot sustain it for 15 minutes gets a capacity ticket against the owning team.",
|
||||
"Stateful services are scaled by hand because adding a node triggers data rebalancing. The owner schedules the change outside peak hours, watches rebalancing progress on the storage dashboard and does not start a second change until the first one has finished.")),
|
||||
new Section("7. Databases", List.of(
|
||||
"The standard database is PostgreSQL 16 with one primary and two replicas in separate availability zones. Replication lag above 15 seconds raises a warning and lag above 60 seconds pages the on-call engineer, because a replica that far behind cannot be promoted safely.",
|
||||
"Automatic failover must complete within 90 seconds. Services reach the database through a connection pooler, so they reconnect on their own, and each service may open at most 200 connections in total. A service that needs more must use the pooler's transaction mode instead of raising the limit.",
|
||||
"Maintenance tasks such as vacuum and index rebuilds run only in the window from 02:00 to 04:00 UTC. Schema changes that need an exclusive lock are not allowed in production at all; the owner must use the online migration tool, which copies the table in the background and swaps it atomically.")),
|
||||
new Section("8. Networking and Certificates", List.of(
|
||||
"All workloads run in private subnets. Outbound traffic leaves through NAT gateways with fixed addresses so that partners can allow-list them, and inbound traffic arrives only through the shared load balancers. The minimum accepted protocol version is TLS 1.3 for every internal and external endpoint.",
|
||||
"Certificates are renewed automatically 30 days before they expire. If a renewal fails, an alert goes to the platform team on the first day and pages them when fewer than 10 days remain, because an expired certificate is the most common cause of a self-inflicted outage.",
|
||||
"DNS records used for failover have a time to live of 60 seconds. All other records use 3600 seconds. Changing a TTL ahead of a planned migration must be done at least one full old-TTL period before the migration starts, so that resolvers have picked up the shorter value.")),
|
||||
new Section("9. Logging and Monitoring", List.of(
|
||||
"Application logs are searchable for 90 days in the hot tier and then archived for a further 400 days in object storage, where they can be restored on request. Logs that contain personal data are redacted at the agent before they leave the host.",
|
||||
"Metrics are stored at 10 second resolution for 7 days, then downsampled to 1 minute resolution for 13 months. Dashboards belong to the service teams, and a dashboard that nobody has opened for 6 months is archived.",
|
||||
"To keep paging humane, each on-call shift has a noise budget of no more than 5 pages. A team that exceeds the budget two shifts in a row must spend the next sprint reducing alert noise before taking new feature work.")),
|
||||
new Section("10. Cost Management", List.of(
|
||||
"Finance and engineering review cloud spend together on the third working day of each month. Any service whose spend is more than 20 percent above its forecast must present a short explanation and a plan at that review.",
|
||||
"Resources without an owner tag are shut down after 14 days of being idle and deleted after another 14 days. The tagging policy is enforced at provisioning time, so an untagged resource can only exist if someone bypassed the pipeline.",
|
||||
"The target for committed-use discounts is to cover 70 percent of steady-state compute. Coverage above 85 percent is discouraged because it removes the ability to shrink the fleet when demand falls.")),
|
||||
new Section("11. Compliance and Vendors", List.of(
|
||||
"The platform is audited for SOC 2 Type II every October, and the auditors receive read-only access to the evidence repository for the whole of that month. Penetration tests by an external firm are carried out twice a year, and every finding rated high or critical must be fixed within 30 days.",
|
||||
"Customers on the Regulated tier have their data stored and processed only in EU regions. A new subprocessor may not receive customer data until a data processing agreement is signed, and the public subprocessor list is updated within 10 business days of any change.",
|
||||
"Employees complete the security awareness course in their first week and again every 12 months. Completion is tracked in the learning system, and access to production is suspended for anyone more than 30 days overdue.")),
|
||||
new Section("12. Decommissioning", List.of(
|
||||
"A service can be deprecated only after its owner has announced the retirement date at least 180 days in advance and has contacted every known consumer directly. The platform team keeps a calendar of retirement dates so that two large migrations do not land in the same month.",
|
||||
"After a customer contract ends, the customer's data is deleted within 30 days and the deletion is confirmed to the customer in writing. Backups that contain the data age out under the normal retention schedule and are not restored for any reason in the meantime.",
|
||||
"When a service is switched off, its DNS records are removed after 14 days of zero traffic, and its credentials are revoked on the same day it stops serving. A final check confirms that no alerts, dashboards or scheduled jobs still reference the retired service.")));
|
||||
|
||||
public static final List<Question> QUESTIONS = List.of(
|
||||
new Question("How long do we keep the daily database snapshots before they are deleted?", "kept for 35 days"),
|
||||
new Question("How often is a restore test performed?", "first Monday of every quarter"),
|
||||
new Question("How far back can a database be recovered to a chosen moment?", "last 7 days"),
|
||||
new Question("How quickly does someone have to start working on the most serious kind of incident?", "within 10 minutes of the page"),
|
||||
new Question("By when must a post-incident review be written?", "within 5 business days"),
|
||||
new Question("How frequently are customers updated on the status page during a major outage?", "every 30 minutes until resolved"),
|
||||
new Question("At what time does the on-call duty change hands each week?", "every Thursday at 10:00 UTC"),
|
||||
new Question("What extra money does an on-call engineer receive?", "120 euros"),
|
||||
new Question("What happens if nobody answers a page for five minutes?", "does not acknowledge a page within 5 minutes"),
|
||||
new Question("On which days and times are releases to production not permitted?", "every Friday after 14:00 UTC"),
|
||||
new Question("What proportion of users sees a new release first, and for how long?", "5 percent of traffic for 20 minutes"),
|
||||
new Question("When does a release get undone without human involvement?", "rises more than 2 percent above the baseline"),
|
||||
new Question("How long does a production login certificate last?", "valid for 8 hours"),
|
||||
new Question("What is the consequence of using the emergency account?", "every use pages the security team"),
|
||||
new Question("After how long does the emergency account's token stop working?", "expires after 60 minutes"),
|
||||
new Question("What CPU load do we aim for when scaling services automatically?", "target of 65 percent"),
|
||||
new Question("How much spare capacity does each region hold?", "reserved headroom of 30 percent"),
|
||||
new Question("How long does the autoscaler wait after removing instances?", "waits 600 seconds"),
|
||||
new Question("When does database replication delay start to raise a warning?", "lag above 15 seconds raises a warning"),
|
||||
new Question("How quickly must a database failover complete?", "within 90 seconds"),
|
||||
new Question("What is the cap on database connections for one service?", "at most 200 connections"),
|
||||
new Question("How long before expiry are TLS certificates replaced?", "renewed automatically 30 days before"),
|
||||
new Question("How long can resolvers cache DNS answers for failover names?", "time to live of 60 seconds"),
|
||||
new Question("What is the oldest TLS version we accept?", "TLS 1.3 for every"),
|
||||
new Question("For how long can I search recent application logs?", "searchable for 90 days"),
|
||||
new Question("What is the page limit for a single on-call shift?", "no more than 5 pages"),
|
||||
new Question("How fine-grained are metrics during their first week?", "10 second resolution"),
|
||||
new Question("When do finance and engineering meet to discuss cloud spend?", "third working day of each month"),
|
||||
new Question("How far above forecast does spending have to go before a team must explain it?", "more than 20 percent above its forecast"),
|
||||
new Question("How much of our steady compute should be covered by committed-use discounts?", "cover 70 percent of steady-state compute"),
|
||||
new Question("In which month does the SOC 2 audit take place?", "SOC 2 Type II every October"),
|
||||
new Question("What is the deadline for fixing serious penetration test findings?", "fixed within 30 days"),
|
||||
new Question("Where may data of customers on the Regulated tier live?", "only in EU regions"),
|
||||
new Question("How much warning must be given before a service is retired?", "at least 180 days in advance"),
|
||||
new Question("How soon after a contract ends is the customer's data erased?", "deleted within 30 days and"),
|
||||
new Question("When are the DNS records of a switched-off service removed?", "removed after 14 days of zero traffic"),
|
||||
new Question("What kind of new subprocessors need a signed agreement before getting data?", "data processing agreement is signed"));
|
||||
|
||||
/** The whole runbook as one string: title, then each section heading and paragraphs separated by blank lines. */
|
||||
public static String fullText() {
|
||||
StringBuilder sb = new StringBuilder(TITLE).append("\n\n");
|
||||
for (Section s : SECTIONS) {
|
||||
sb.append(s.title()).append("\n\n");
|
||||
for (String p : s.paragraphs()) {
|
||||
sb.append(p).append("\n\n");
|
||||
}
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
|
||||
/** Title, then each section heading followed by its paragraphs, in reading order. */
|
||||
public static List<String> blocks() {
|
||||
java.util.List<String> blocks = new java.util.ArrayList<>();
|
||||
blocks.add(TITLE);
|
||||
for (Section s : SECTIONS) {
|
||||
blocks.add(s.title());
|
||||
blocks.addAll(s.paragraphs());
|
||||
}
|
||||
return blocks;
|
||||
}
|
||||
|
||||
private Corpus() {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentTransformer;
|
||||
|
||||
/**
|
||||
* One chunk per paragraph, each tagged with the section heading it sits under. Line-oriented text (what the
|
||||
* Tika reader returns for a Word file) is read as: a short numbered line with no full stop is a heading, any
|
||||
* other line is a paragraph. With {@code prefixHeading} the heading is also written in front of the paragraph,
|
||||
* so the embedding sees which section the paragraph belongs to.
|
||||
*/
|
||||
public class HeadingAwareChunker implements DocumentTransformer {
|
||||
|
||||
private static final Pattern HEADING = Pattern.compile("^\\d+\\. [^.]{1,60}$");
|
||||
|
||||
private final boolean prefixHeading;
|
||||
|
||||
public HeadingAwareChunker(boolean prefixHeading) {
|
||||
this.prefixHeading = prefixHeading;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> apply(List<Document> documents) {
|
||||
List<Document> out = new ArrayList<>();
|
||||
for (Document doc : documents) {
|
||||
String section = "";
|
||||
int index = 0;
|
||||
for (String raw : doc.getText().split("\\R")) {
|
||||
String line = raw.replaceAll("\\s+", " ").strip();
|
||||
if (line.isEmpty()) {
|
||||
continue;
|
||||
}
|
||||
if (HEADING.matcher(line).matches()) {
|
||||
section = line;
|
||||
continue;
|
||||
}
|
||||
if (section.isEmpty()) {
|
||||
continue; // the document title, before the first section
|
||||
}
|
||||
Map<String, Object> metadata = new LinkedHashMap<>(doc.getMetadata());
|
||||
metadata.put("section", section);
|
||||
metadata.put("chunk_index", index++);
|
||||
out.add(Document.builder().text(prefixHeading ? section + "\n" + line : line).metadata(metadata).build());
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.nio.file.Path;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Set;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentReader;
|
||||
import org.springframework.ai.document.DocumentTransformer;
|
||||
import org.springframework.ai.reader.jsoup.JsoupDocumentReader;
|
||||
import org.springframework.ai.reader.jsoup.config.JsoupDocumentReaderConfig;
|
||||
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
|
||||
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
|
||||
import org.springframework.ai.reader.tika.TikaDocumentReader;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
import org.springframework.core.io.FileSystemResource;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
|
||||
/**
|
||||
* Reads one file, splits it, tags the chunks and writes them to a pgvector table so that running it again
|
||||
* on the same file is harmless:
|
||||
* <ul>
|
||||
* <li>chunks that already exist (same id) are skipped, so nothing is embedded twice;</li>
|
||||
* <li>chunks that exist in the table for this source but not in the new version are deleted;</li>
|
||||
* <li>only the new chunks are embedded and written.</li>
|
||||
* </ul>
|
||||
*/
|
||||
public class IngestionPipeline {
|
||||
|
||||
public record Result(int chunks, int embedded, int unchanged, int deleted) {
|
||||
}
|
||||
|
||||
private final VectorStore store;
|
||||
private final JdbcTemplate jdbc;
|
||||
private final String table;
|
||||
private final DocumentTransformer splitter;
|
||||
|
||||
public IngestionPipeline(VectorStore store, JdbcTemplate jdbc, String table, DocumentTransformer splitter) {
|
||||
this.store = store;
|
||||
this.jdbc = jdbc;
|
||||
this.table = table;
|
||||
this.splitter = splitter;
|
||||
}
|
||||
|
||||
public Result ingest(Path file) {
|
||||
String name = file.getFileName().toString();
|
||||
String format = name.substring(name.lastIndexOf('.') + 1).toLowerCase(Locale.ROOT);
|
||||
|
||||
// 1. read, then join every page or element into one text so chunk boundaries do not depend on the file's layout
|
||||
List<Document> read = reader(file, format).get();
|
||||
String text = read.stream().map(Document::getText).collect(Collectors.joining("\n"));
|
||||
List<Document> whole = List.of(new Document(text));
|
||||
|
||||
// 2. split and tag
|
||||
List<Document> chunks = new MetadataEnricher(name, format).apply(splitter.apply(whole));
|
||||
if (chunks.isEmpty() && !text.isBlank()) {
|
||||
// a splitter that does not understand the reader's text shape returns nothing; say so instead of writing zero rows
|
||||
throw new IllegalStateException(name + ": " + text.length() + " characters read, 0 chunks produced");
|
||||
}
|
||||
|
||||
// 3. work out what is already stored for this source
|
||||
Set<String> existing = new HashSet<>(jdbc.queryForList(
|
||||
"select id::text from " + table + " where metadata->>'source' = ?", String.class, name));
|
||||
Set<String> wanted = chunks.stream().map(Document::getId).collect(Collectors.toSet());
|
||||
|
||||
List<String> stale = existing.stream().filter(id -> !wanted.contains(id)).toList();
|
||||
if (!stale.isEmpty()) {
|
||||
store.delete(stale);
|
||||
}
|
||||
List<Document> fresh = chunks.stream().filter(c -> !existing.contains(c.getId())).toList();
|
||||
if (!fresh.isEmpty()) {
|
||||
store.add(fresh);
|
||||
}
|
||||
return new Result(chunks.size(), fresh.size(), chunks.size() - fresh.size(), stale.size());
|
||||
}
|
||||
|
||||
static DocumentReader reader(Path file, String format) {
|
||||
FileSystemResource resource = new FileSystemResource(file);
|
||||
return switch (format) {
|
||||
case "pdf" -> new PagePdfDocumentReader(resource, PdfDocumentReaderConfig.builder().withPageBottomMargin(60).build());
|
||||
// headings and paragraphs as separate lines: JsoupDocumentReader joins elements with the separator, and a line-oriented chunker needs the lines
|
||||
case "html", "htm" -> new JsoupDocumentReader(resource, JsoupDocumentReaderConfig.builder()
|
||||
.selector("article h1, article h2, article p").separator("\n").build());
|
||||
default -> new TikaDocumentReader(resource);
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.security.MessageDigest;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HexFormat;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentTransformer;
|
||||
|
||||
/**
|
||||
* Gives every chunk the same metadata whatever reader produced it: {@code source}, {@code format},
|
||||
* {@code chunk_index} and a {@code content_hash}, and a deterministic id derived from them. The id is the
|
||||
* whole idempotency story: the same chunk of the same file always gets the same id, so writing it twice
|
||||
* is an upsert, and a chunk whose text changed gets a new id.
|
||||
*/
|
||||
public class MetadataEnricher implements DocumentTransformer {
|
||||
|
||||
private final String source;
|
||||
private final String format;
|
||||
|
||||
public MetadataEnricher(String source, String format) {
|
||||
this.source = source;
|
||||
this.format = format;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> apply(List<Document> chunks) {
|
||||
List<Document> out = new ArrayList<>();
|
||||
for (int i = 0; i < chunks.size(); i++) {
|
||||
Document chunk = chunks.get(i);
|
||||
String hash = sha256(chunk.getText()).substring(0, 16);
|
||||
Map<String, Object> metadata = new LinkedHashMap<>(chunk.getMetadata());
|
||||
metadata.put("source", source);
|
||||
metadata.put("format", format);
|
||||
metadata.put("chunk_index", i);
|
||||
metadata.put("content_hash", hash);
|
||||
String id = UUID.nameUUIDFromBytes((source + "#" + hash).getBytes(StandardCharsets.UTF_8)).toString();
|
||||
out.add(Document.builder().id(id).text(chunk.getText()).metadata(metadata).build());
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
static String sha256(String text) {
|
||||
try {
|
||||
return HexFormat.of().formatHex(MessageDigest.getInstance("SHA-256").digest(text.getBytes(StandardCharsets.UTF_8)));
|
||||
} catch (NoSuchAlgorithmException e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.document.MetadataMode;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.transformers.TransformersEmbeddingModel;
|
||||
|
||||
/**
|
||||
* all-MiniLM-L6-v2 (384 dimensions) running inside the JVM through ONNX Runtime. Both model files come from
|
||||
* Hugging Face and are cached under the user's home directory, so the first run downloads about 90 MB and
|
||||
* later runs are offline. No API key is involved, which is what lets anyone repeat the measurements.
|
||||
*/
|
||||
public final class MiniLm {
|
||||
|
||||
private static final String HF = "https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2/resolve/main/";
|
||||
private static EmbeddingModel shared;
|
||||
|
||||
private MiniLm() {
|
||||
}
|
||||
|
||||
/** The shared instance, which embeds the text of a Document only (MetadataMode.NONE, the constructor default). */
|
||||
public static synchronized EmbeddingModel get() {
|
||||
if (shared == null) {
|
||||
shared = create(MetadataMode.NONE);
|
||||
}
|
||||
return shared;
|
||||
}
|
||||
|
||||
/** A fresh instance whose Document embedding content is built with the given metadata mode. */
|
||||
public static EmbeddingModel create(MetadataMode mode) {
|
||||
try {
|
||||
TransformersEmbeddingModel m = new TransformersEmbeddingModel(mode);
|
||||
m.setTokenizerResource(HF + "tokenizer.json");
|
||||
m.setModelResource(HF + "onnx/model.onnx");
|
||||
m.setResourceCacheDirectory(System.getProperty("user.home") + "/.cache/minilm");
|
||||
m.setTokenizerOptions(Map.of("padding", "true", "truncation", "true", "maxLength", "256"));
|
||||
m.afterPropertiesSet();
|
||||
return m;
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.document.DocumentTransformer;
|
||||
|
||||
/**
|
||||
* Groups whole sentences into chunks of about {@code targetTokens}, repeating the last
|
||||
* {@code overlapSentences} sentences at the start of the next chunk. Unlike TokenTextSplitter it never cuts
|
||||
* inside a sentence, so an answer that sits in one sentence cannot be torn in half by a chunk boundary.
|
||||
*/
|
||||
public class SentenceSplitter implements DocumentTransformer {
|
||||
|
||||
/** A full stop after a letter, followed by a capital: "UTC. Daily" splits, "1. Backups" does not. */
|
||||
private static final String SENTENCE_END = "(?<=[a-z\\)][.!?])\\s+(?=[A-Z])";
|
||||
|
||||
private final int targetTokens;
|
||||
private final int overlapSentences;
|
||||
|
||||
public SentenceSplitter(int targetTokens, int overlapSentences) {
|
||||
this.targetTokens = targetTokens;
|
||||
this.overlapSentences = overlapSentences;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<Document> apply(List<Document> documents) {
|
||||
List<Document> out = new ArrayList<>();
|
||||
for (Document doc : documents) {
|
||||
List<String> sentences = new ArrayList<>();
|
||||
for (String line : doc.getText().split("\\R")) {
|
||||
String clean = line.replaceAll("\\s+", " ").strip();
|
||||
if (!clean.isEmpty()) {
|
||||
sentences.addAll(List.of(clean.split(SENTENCE_END)));
|
||||
}
|
||||
}
|
||||
List<String> window = new ArrayList<>();
|
||||
int tokens = 0;
|
||||
int index = 0;
|
||||
for (String sentence : sentences) {
|
||||
int t = Tokens.count(sentence);
|
||||
if (!window.isEmpty() && tokens + t > targetTokens) {
|
||||
out.add(chunk(doc, window, index++));
|
||||
List<String> keep = overlapSentences == 0 ? List.of()
|
||||
: window.subList(Math.max(0, window.size() - overlapSentences), window.size());
|
||||
window = new ArrayList<>(keep);
|
||||
tokens = window.stream().mapToInt(Tokens::count).sum();
|
||||
}
|
||||
window.add(sentence);
|
||||
tokens += t;
|
||||
}
|
||||
if (!window.isEmpty()) {
|
||||
out.add(chunk(doc, window, index));
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
private static Document chunk(Document parent, List<String> sentences, int index) {
|
||||
Map<String, Object> metadata = new LinkedHashMap<>(parent.getMetadata());
|
||||
metadata.put("chunk_index", index);
|
||||
return Document.builder().text(String.join(" ", sentences)).metadata(metadata).build();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
package com.ankurm.ragingestion;
|
||||
|
||||
import com.knuddels.jtokkit.Encodings;
|
||||
import com.knuddels.jtokkit.api.Encoding;
|
||||
import com.knuddels.jtokkit.api.EncodingType;
|
||||
|
||||
/** Token counts in the same encoding TokenTextSplitter uses by default (cl100k_base). */
|
||||
public final class Tokens {
|
||||
|
||||
private static final Encoding ENCODING = Encodings.newDefaultEncodingRegistry().getEncoding(EncodingType.CL100K_BASE);
|
||||
|
||||
private Tokens() {
|
||||
}
|
||||
|
||||
public static int count(String text) {
|
||||
return ENCODING.countTokens(text);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user