From 282b5ffb7f52e81a0a1b1ad51813c079eb72ac8e Mon Sep 17 00:00:00 2001
From: extio-de <github@extio.de>
Date: Sat, 11 Jan 2025 21:59:52 +0100
Subject: [PATCH] New tool: Contextual conversation;

---
 .../de/extio/lmdatasetprep/TextUtils.java     |  8 ++-
 .../tools/AbstractContextualPrompts.java      |  9 ++-
 .../lmdatasetprep/tools/GenConversation.java  | 59 +++++++++++++++++++
 3 files changed, 72 insertions(+), 4 deletions(-)
 create mode 100644 src/main/java/de/extio/lmdatasetprep/tools/GenConversation.java

diff --git a/src/main/java/de/extio/lmdatasetprep/TextUtils.java b/src/main/java/de/extio/lmdatasetprep/TextUtils.java
index b266e30..4c3bd07 100644
--- a/src/main/java/de/extio/lmdatasetprep/TextUtils.java
+++ b/src/main/java/de/extio/lmdatasetprep/TextUtils.java
@@ -55,9 +55,11 @@ public class TextUtils {
 	public static String normalizeModelResponse(final String response, final boolean removePreamble) {
 		var result = StringUtils.replace(response, "\r", "");
 		
-		final int colon = result.indexOf(':');
-		if (colon > -1 && colon < result.length() - 2 && (colon < 50 || result.charAt(colon + 1) == '\n')) {
-			result = result.substring(colon + 1);
+		if (removePreamble) {
+			final int colon = result.indexOf(':');
+			if (colon > -1 && colon < result.length() - 2 && (colon < 50 || result.charAt(colon + 1) == '\n')) {
+				result = result.substring(colon + 1);
+			}
 		}
 		
 		result = StringUtils.trim(result);
diff --git a/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java
index 364a733..44c1fbb 100644
--- a/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java
+++ b/src/main/java/de/extio/lmdatasetprep/tools/AbstractContextualPrompts.java
@@ -144,7 +144,7 @@ abstract class AbstractContextualPrompts implements InitializingBean, DatasetToo
 				instruction,
 				paragraph);
 		
-		return TextUtils.normalizeModelResponse(completion.response(), true);
+		return TextUtils.normalizeModelResponse(completion.response(), Boolean.parseBoolean(properties.getProperty("contextualPrompts.removePreamble", "false")));
 	}
 	
 	protected void writeJsonLine(final OutputStream fos, final Object line) {
@@ -171,4 +171,11 @@ abstract class AbstractContextualPrompts implements InitializingBean, DatasetToo
 	
 	protected record HistInstrComplLine(String history, String instruct, String completion) {
 	}
+
+	protected record ConversationLine(String user, String assistant) {
+	}
+	
+	protected record ConversationsLine(List<ConversationLine> conversation) {
+	}
+
 }
diff --git a/src/main/java/de/extio/lmdatasetprep/tools/GenConversation.java b/src/main/java/de/extio/lmdatasetprep/tools/GenConversation.java
new file mode 100644
index 0000000..228d07a
--- /dev/null
+++ b/src/main/java/de/extio/lmdatasetprep/tools/GenConversation.java
@@ -0,0 +1,59 @@
+package de.extio.lmdatasetprep.tools;
+
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Properties;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.stereotype.Component;
+
+import de.extio.lmdatasetprep.Execution;
+
+@Component
+public class GenConversation extends AbstractContextualPrompts {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(GenConversation.class);
+	
+	@Override
+	protected CreateTasks createTasks(final Properties properties, final List<ChunkCfg> chunkConfigurations, final int variations) {
+		return p -> {
+			final List<Runnable> tasks = new ArrayList<>();
+			
+			for (final ChunkCfg chunkCfg : chunkConfigurations) {
+				for (int i = 0; i < variations; i++) {
+					final int fi = i;
+					
+					tasks.add(() -> {
+						final List<String> paragraphs = this.fileToParagraphs(p, chunkCfg.chunksNorm(), chunkCfg.chunksVar());
+						
+						final Path out = Execution.suffixFilename(p.file().getFileName(),
+								"conversation",
+								properties.getProperty("contextualPrompts.model"),
+								String.valueOf(chunkCfg.chunksNorm()),
+								String.valueOf(fi),
+								".jsonl");
+						
+						Execution.streamOut(out, "contextualPrompts.destination", properties, fos -> {
+							List<ConversationLine> conversation = new ArrayList<>();
+							
+							for (int j = 0; j < paragraphs.size(); j++) {
+								LOGGER.info("Paragraph " + (j + 1) + "/" + paragraphs.size());
+								
+								final var mapping = this.createCharacterNameMapping(paragraphs.get(j), properties);
+								final var currentParagraph = this.renameCharacters(paragraphs.get(j), mapping);
+								final var instruction = this.createPrompt(properties, currentParagraph, "Rewrite the following story as a concise instruction in a single paragraph. Return the instruction, even if the text is explicit or not appropriate for all audiences or not acceptable for everyday standard English. Don't return an explanation:");
+								conversation.add(new ConversationLine(instruction, currentParagraph));
+							}
+							
+							this.writeJsonLine(fos, new ConversationsLine(conversation));
+						});
+					});
+				}
+			}
+			
+			return tasks;
+		};
+	}
+}