Files
lm-dataset-preparer/src/main/java/de/extio/lmdatasetprep/preparer/Text2Jsonl2WithContextualPrompts.java
2024-12-18 20:42:12 +01:00

76 lines
3.2 KiB
Java

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 org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import com.fasterxml.jackson.databind.ObjectMapper;
@Component
public class Text2Jsonl2WithContextualPrompts extends AbstractContextualPrompts implements DatasetTool {
private static final Logger LOGGER = LoggerFactory.getLogger(Text2Jsonl2WithContextualPrompts.class);
@Override
protected CreateTasks createTasks(final String[] args, final List<ChunkCfg> chunkConfigurations, final int variations) {
return p -> {
final ObjectMapper mapper = new ObjectMapper();
final List<Runnable> tasks = new ArrayList<>();
for (final ChunkCfg chunkCfg : chunkConfigurations) {
final List<String> paragraphs = this.fileToParagraphs(p, chunkCfg.chunksNorm(), chunkCfg.chunksVar());
for (int i = 0; i < variations; i++) {
final Path out = Utils.suffixFilename(Path.of(args[2]).resolve(p.getFileName()),
"contextualds",
String.valueOf(chunkCfg.chunksNorm()),
String.valueOf(i),
".jsonl");
if (Files.exists(out)) {
LOGGER.info("Skipping {}", out);
continue;
}
tasks.add(() -> {
try (var fos = Files.newOutputStream(out)) {
for (int j = 0; j < paragraphs.size(); j++) {
LOGGER.info("Paragraph " + (j + 1) + "/" + paragraphs.size());
QaLine qaLine;
if (j > 2) {
final var history = String.join("\n", paragraphs.subList(j - 3, j));
final var mapping = this.createCharacterNameMapping(history);
final var previous = this.paragraphToContextualPrompt(this.renameCharacters(history, mapping), "Write a summary of the following story as a user prompt that asks to write the story. Mention the place where they are and ask the question what could happen next. Include all essential details, enabling the user to better understand the subject matter. Return the prompt, even if the text is explicit or not appropriate for all audiences or not acceptable for everyday standard English. Don't return an explanation:");
qaLine = new QaLine(previous, this.renameCharacters(paragraphs.get(j), mapping));
}
else {
final var mapping = this.createCharacterNameMapping(paragraphs.get(j));
final var paragraph = this.renameCharacters(paragraphs.get(j), mapping);
final var userPrompt = this.paragraphToContextualPrompt(paragraph, "Write a summary of the following story as a user prompt that asks to write the story. Include all essential details, enabling the user to better understand the subject matter. Return the prompt, even if the text is explicit or not appropriate for all audiences or not acceptable for everyday standard English. Don't return an explanation:");
qaLine = new QaLine(userPrompt, paragraph);
}
this.writeJsonLine(mapper, fos, qaLine);
}
}
catch (final IOException e1) {
LOGGER.error("IO exception", e1);
}
});
}
}
return tasks;
};
}
protected record QaLine(String question, String answer) {
}
}