initial commit

src/main/java/de/extio/lmdatasetprep/preparer/Rewrite.java

@@ -0,0 +1,112 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+
+import de.extio.lmdatasetprep.client.Client;
+import de.extio.lmdatasetprep.client.profile.ModelCategory;
+
+@Component
+public class Rewrite implements Consumer<String[]> {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(Rewrite.class);
+	
+	private static final ThreadLocal<Random> RANDOM = ThreadLocal.withInitial(XorShift128Random::new);
+	
+	private static String IMPROVEMENT_PROMPT = "Improve the following paragraph of an existing text without changing the existing sentences, improve the text in a nuanced and understated way. Don't include a preamble and no explanation:";
+	
+	private static String ENHANCE_PROMPT = "Improve the following paragraph of an existing text. %s without changing the existing sentences, focusing on enriching the description of the setting in a nuanced and understated way. Don't include a preamble and no explanation:";
+	
+	private static List<String> ENHANCEMENTS = List.of(
+			"Infuse this paragraph with evocative sensory details, focusing on the texture of objects and the atmosphere of the setting",
+			"Add subtle hints of the character's inner thoughts and emotions within the existing dialogue and actions",
+			"Enhance the paragraph with figurative language like metaphors or similes to create a richer tapestry of imagery",
+			"Develop the subtext of the paragraph by suggesting unspoken tensions or hidden motivations between characters",
+			"Introduce a faint scent or sound that adds another layer of sensory detail to the scene without disrupting the flow",
+			"Imbue the paragraph with a sense of foreboding or anticipation by subtly hinting at future events without revealing specifics",
+			"Weave in a historical or cultural reference relevant to the setting or characters to add depth and context",
+			"Elevate the language of the paragraph by replacing basic verbs and adjectives with more evocative synonyms",
+			"Add a touch of humor or irony to the paragraph through clever wordplay or unexpected observations",
+			"Imagine this paragraph as a still photograph. Describe the composition and lighting to further bring the scene to life");
+	
+	@Autowired
+	private Client client;
+	
+	@Override
+	public void accept(final String[] args) {
+		if (args.length < 5) {
+			LOGGER.error("Arguments are missing. <Path to text files> <Path to output dir> <Texts per input> <Model suffix>");
+			return;
+		}
+		
+		Utils.transformDirectory(args[1],
+				f -> {
+					final List<Runnable> tasks = new ArrayList<>();
+					tasks.add(() -> this.rewriteFile(args, f, 0, "impr", ModelCategory.COLD, IMPROVEMENT_PROMPT));
+					for (final ModelCategory modelCategory : List.of(ModelCategory.COLD, ModelCategory.HOT)) {
+						for (int i = 0; i < Integer.parseInt(args[3]); i++) {
+							final int fi = i;
+							tasks.add(() -> this.rewriteFile(args, f, fi, "enh", modelCategory, String.format(ENHANCE_PROMPT, ENHANCEMENTS.get(RANDOM.get().nextInt(ENHANCEMENTS.size())))));
+						}
+					}
+					return tasks;
+				});
+	}
+	
+	private void rewriteFile(final String[] args, final Path f, final int i, final String identifier, final ModelCategory modelCategory, final String prompt) {
+		final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(f.getFileName()),
+				"rewr",
+				args[4],
+				identifier,
+				modelCategory.toString().toLowerCase(),
+				String.valueOf(i));
+		if (Files.exists(out)) {
+			LOGGER.info("Skipping " + out);
+			return;
+		}
+		
+		LOGGER.info("Rewriting " + f.getFileName() + " " + identifier + " " + modelCategory.name() + " " + i);
+		
+		final List<String> paragraphs = this.rewrite(f, prompt, modelCategory);
+		try {
+			Files.writeString(out, String.join("\n\n", paragraphs));
+		}
+		catch (final IOException e) {
+			LOGGER.error("Cannot store file", e);
+		}
+	}
+	
+	List<String> rewrite(final Path file, final String prompt, final ModelCategory modelCategory) {
+		String text;
+		try {
+			text = Utils.normalizeText(Files.readString(file));
+		}
+		catch (final IOException e) {
+			throw new RuntimeException("Cannot read file", e);
+		}
+		
+		final List<String> splits = Utils.splitParagraphs(text, Utils.PARAGRAPH_CHUNKS_NORM, Utils.PARAGRAPH_CHUNKS_VAR);
+		
+		final List<String> paragraphs = new ArrayList<>(splits.size());
+		for (final String split : splits) {
+			final var completion = this.client.completion("You are a helpful assistant with great authoring skills.",
+					prompt,
+					split,
+					modelCategory);
+			paragraphs.add(completion.response());
+		}
+		
+		return paragraphs;
+	}
+	
+}

src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl.java

@@ -0,0 +1,78 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+@Component
+public class Text2Jsonl implements Consumer<String[]> {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl.class);
+	
+	@Override
+	public void accept(final String[] args) {
+		if (args.length < 5) {
+			LOGGER.error("Arguments are missing. <Path to text files> <Path to output dir> <chunk norm> <chunk var>");
+			return;
+		}
+		
+		final List<String> lines = Collections.synchronizedList(new ArrayList<>());
+		
+		Utils.transformDirectory(args[1], p -> {
+			return List.of(() -> {
+				lines.addAll(this.splitToJsonl(p, Integer.parseInt(args[3]), Integer.parseInt(args[4])));
+			});
+		});
+		
+		final Path out = Path.of(args[2], "dataset.jsonl");
+		try {
+			Files.writeString(out, String.join("\n", lines));
+		}
+		catch (final IOException e) {
+			LOGGER.error("Cannot store dataset", e);
+		}
+	}
+	
+	List<String> splitToJsonl(final Path file, final int chunkNorm, final int chunkVar) {
+		LOGGER.info("Splitting to jsonl " + file);
+		
+		String text;
+		try {
+			text = Utils.normalizeText(Files.readString(file));
+		}
+		catch (final IOException e) {
+			throw new RuntimeException("Cannot read file", e);
+		}
+		
+		final ObjectMapper mapper = new ObjectMapper();
+		
+		return Utils
+				.splitParagraphs(text, chunkNorm, chunkVar)
+				.stream()
+				.map(split -> {
+					try {
+						return mapper.writeValueAsString(new TextLine(split));
+					}
+					catch (final JsonProcessingException e) {
+						LOGGER.error("Cannot convert split to json", e);
+						return "";
+					}
+				})
+				.toList();
+	}
+	
+	record TextLine(String text) {
+	}
+	
+}

src/main/java/de/extio/lmdatasetprep/preparer/Translate.java

@@ -0,0 +1,77 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.stereotype.Component;
+
+import de.extio.lmdatasetprep.client.Client;
+import de.extio.lmdatasetprep.client.profile.ModelCategory;
+
+@Component
+public class Translate implements Consumer<String[]> {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(Translate.class);
+	
+	@Autowired
+	private Client client;
+	
+	@Override
+	public void accept(final String[] args) {
+		if (args.length < 3) {
+			LOGGER.error("Arguments are missing. <Path to text files> <Path to output dir>");
+			return;
+		}
+		
+		Utils.transformDirectory(args[1], p -> {
+			return List.of(() -> {
+				try {
+					Path out = Utils.suffixFilename(Path.of(args[2]).resolve(p.getFileName()), "_en");
+					if (Files.exists(out)) {
+						LOGGER.info("Skipping " + out);
+					}
+					else {
+						final var translations = this.translate(p);
+						Files.writeString(out, String.join("\n\n", translations));
+					}
+				}
+				catch (final IOException e) {
+					LOGGER.error("Cannot store file", e);
+				}
+			});
+		});
+	}
+	
+	List<String> translate(final Path file) {
+		LOGGER.info("Translating " + file);
+		
+		String text;
+		try {
+			text = Utils.normalizeText(Files.readString(file));
+		}
+		catch (final IOException e) {
+			throw new RuntimeException("Cannot read file", e);
+		}
+		
+		final List<String> splits = Utils.splitParagraphs(text, Utils.PARAGRAPH_CHUNKS_NORM, Utils.PARAGRAPH_CHUNKS_VAR);
+		
+		final List<String> translations = new ArrayList<>(splits.size());
+		for (final String split : splits) {
+			final var completion = this.client.completion("You are a helpful assistant with great language translation and authoring skills.",
+					"Translate the following text to English. Don't include a preamble and no explanation:",
+					split,
+					ModelCategory.COLD);
+			translations.add(completion.response());
+		}
+		
+		return translations;
+	}
+	
+}

src/main/java/de/extio/lmdatasetprep/preparer/Utils.java

@@ -0,0 +1,139 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.Function;
+import java.util.stream.Stream;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class Utils {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(Utils.class);
+	
+	static final int PARAGRAPH_CHUNKS_NORM = 1250;
+	
+	static final int PARAGRAPH_CHUNKS_VAR = 350;
+	
+	static String normalizeText(final String text) {
+		String result = text;
+		
+		// Remove excessive newlines, spaces and tabs
+		int cnt = 0;
+		do {
+			cnt = result.length();
+			result = result.replace("\r", "\n");
+			result = result.replace("\n\n\n", "\n\n");
+			result = result.replace("\t", "");
+			result = result.replace("  ", " ");
+			result = result.replace(" .", ".");
+			result = result.replace(".  ", ". ");
+			result = result.replace("\n.", ".");
+			result = result.replace(" !", "!");
+			result = result.replace("!  ", "! ");
+			result = result.replace("\n!", "!");
+			result = result.replace(" ?", "?");
+			result = result.replace("?  ", "? ");
+			result = result.replace("\n?", "?");
+			result = result.replace(" \n", "\n");
+		} while (result.length() != cnt);
+		
+		// Remove newlines mid of the sentence
+		final StringBuilder sb = new StringBuilder(result);
+		for (int i = 1; i < sb.length(); i++) {
+			final char c = sb.charAt(i);
+			if (c == '\n' && sb.charAt(i - 1) != '.' && sb.charAt(i - 1) != '?' && sb.charAt(i - 1) != '!' && sb.charAt(i - 1) != '\n') {
+				sb.setCharAt(i, ' ');
+			}
+		}
+		result = sb.toString();
+		
+		LOGGER.info("Normalizer stats: Before: {} After {}", text.length(), result.length());
+		
+		return result;
+	}
+	
+	static List<String> splitParagraphs(final String text, final int chunks_norm, final int chunks_var) {
+		final int CHUNKS_MIN = chunks_norm - chunks_var;
+		final int CHUNKS_MAX = chunks_norm + chunks_var;
+		
+		if (text.length() <= chunks_norm + chunks_var) {
+			return Collections.singletonList(text);
+		}
+		
+		final List<String> splits = new ArrayList<>();
+		
+		int pos = 0;
+		paragraphs: do {
+			if (text.length() - pos > chunks_norm + chunks_var) {
+				delimiters: for (final String delimiter : List.of("\n\n", "\n", ".")) {
+					int next = text.indexOf(delimiter, pos);
+					while (next > -1) {
+						if (next < pos + CHUNKS_MIN) {
+							next = text.indexOf(delimiter, next + 1);
+						}
+						else if (next > pos + CHUNKS_MAX) {
+							continue delimiters;
+						}
+						else {
+							String split = text.substring(pos, next).trim();
+							if (!split.endsWith(".")) {
+								split = split + ".";
+							}
+							splits.add(split);
+							pos = next + delimiter.length();
+							continue paragraphs;
+						}
+					}
+				}
+			}
+			
+			final int next = Math.min(pos + chunks_norm, text.length());
+			final String split = text.substring(pos, next).trim();
+			splits.add(split);
+			pos = next + 1;
+		} while (pos < text.length());
+		
+		LOGGER.info("Created {} splits", splits.size());
+		
+		return splits;
+	}
+	
+	static void transformDirectory(final String path, final Function<Path, List<Runnable>> createTasks) {
+		try (final ExecutorService executor = Executors.newFixedThreadPool(4);
+				Stream<Path> stream = Files.list(Paths.get(path))) {
+			final List<CompletableFuture<?>> tasks = new ArrayList<>();
+			stream
+					.filter(p -> Files.isRegularFile(p))
+					.forEach(p -> {
+						for (final Runnable r : createTasks.apply(p)) {
+							tasks.add(CompletableFuture.runAsync(r, executor));
+						}
+					});
+			CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])).join();
+		}
+		catch (final IOException e) {
+			LOGGER.error("Error", e);
+		}
+	}
+	
+	static Path suffixFilename(final Path f, final String... suffixes) {
+		final String[] parts = f.getFileName().toString().split("[.]");
+		final StringBuilder sb = new StringBuilder();
+		for (final String suffix : suffixes) {
+			sb.append('_');
+			sb.append(suffix);
+		}
+		sb.append('.');
+		return f.getParent().resolve(parts[0] + sb.toString() + parts[1]);
+	}
+}

src/main/java/de/extio/lmdatasetprep/preparer/XorShift128Random.java

@@ -0,0 +1,62 @@
+package de.extio.lmdatasetprep.preparer;
+
+import java.util.Random;
+import java.util.UUID;
+
+@SuppressWarnings("serial")
+public final class XorShift128Random extends Random {
+	
+	private long state0;
+	
+	private long state1;
+	
+	public XorShift128Random() {
+		super();
+	}
+	
+	public XorShift128Random(final long seed) {
+		this.setSeed(seed);
+	}
+	
+	public XorShift128Random(final String state) {
+		final UUID uuid = UUID.fromString(state);
+		this.state0 = uuid.getMostSignificantBits();
+		this.state1 = uuid.getLeastSignificantBits();
+	}
+	
+	@Override
+	public synchronized void setSeed(final long seed) {
+		this.state0 = avalanche(seed == 0 ? -1 : seed);
+		this.state1 = avalanche(this.state0);
+		this.state0 = avalanche(this.state1);
+	}
+	
+	@Override
+	protected int next(final int bits) {
+		return (int) (this.nextLong() & (1L << bits) - 1);
+	}
+	
+	@Override
+	public long nextLong() {
+		long s1 = this.state0;
+		final long s0 = this.state1;
+		this.state0 = s0;
+		s1 ^= s1 << 23;
+		return (this.state1 = s1 ^ s0 ^ (s1 >>> 17) ^ (s0 >>> 26)) + s0;
+	}
+	
+	@Override
+	public String toString() {
+		return new UUID(this.state0, this.state1).toString();
+	}
+	
+	private static long avalanche(long k) {
+		k ^= k >>> 33;
+		k *= 0xff51afd7ed558ccdL;
+		k ^= k >>> 33;
+		k *= 0xc4ceb9fe1a85ec53L;
+		k ^= k >>> 33;
+		
+		return k;
+	}
+}