Switch dataset preparer to lmlib;
This commit is contained in:
extio-de
@@ -0,0 +1,176 @@
|
||||
package de.extio.lmdatasetprep.preparer;
|
||||
|
||||
import java.io.BufferedReader;
|
||||
import java.io.IOException;
|
||||
import java.io.InputStreamReader;
|
||||
import java.io.OutputStream;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Random;
|
||||
import java.util.concurrent.ThreadLocalRandom;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.beans.factory.InitializingBean;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.core.io.ClassPathResource;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
|
||||
import de.extio.lmlib.client.ClientService;
|
||||
import de.extio.lmlib.profile.ModelCategory;
|
||||
|
||||
abstract class AbstractContextualPrompts implements InitializingBean, Consumer<String[]> {
|
||||
|
||||
private static final Logger LOGGER = LoggerFactory.getLogger(AbstractContextualPrompts.class);
|
||||
|
||||
private static final byte[] NEWLINE = "\n".getBytes(StandardCharsets.UTF_8);
|
||||
|
||||
@Autowired
|
||||
protected ClientService clientService;
|
||||
|
||||
private List<String> maleNames = new ArrayList<>();
|
||||
|
||||
private List<String> femaleNames = new ArrayList<>();
|
||||
|
||||
@Override
|
||||
public void afterPropertiesSet() throws Exception {
|
||||
this.maleNames = new BufferedReader(new InputStreamReader(new ClassPathResource("male.txt").getInputStream())).lines().collect(Collectors.toList());
|
||||
this.femaleNames = new BufferedReader(new InputStreamReader(new ClassPathResource("female.txt").getInputStream())).lines().collect(Collectors.toList());
|
||||
}
|
||||
|
||||
@Override
|
||||
public void accept(final String[] args) {
|
||||
if (args.length < 6) {
|
||||
LOGGER.error("Arguments are missing. <Path to text files> <Path to output dir> <variations> [<chunk norm> <chunk var>, ...]");
|
||||
return;
|
||||
}
|
||||
final int variations = Integer.parseInt(args[3]);
|
||||
final List<ChunkCfg> chunkConfigurations = new ArrayList<>();
|
||||
for (int i = 4; i < args.length; i += 2) {
|
||||
final var cfg = new ChunkCfg(Integer.parseInt(args[i]), Integer.parseInt(args[i + 1]));
|
||||
chunkConfigurations.add(cfg);
|
||||
}
|
||||
|
||||
Utils.transformDirectory(args[1], this.createTasks(args, chunkConfigurations, variations)::createTasks);
|
||||
}
|
||||
|
||||
@FunctionalInterface
|
||||
static interface CreateTasks {
|
||||
|
||||
List<Runnable> createTasks(Path path);
|
||||
|
||||
}
|
||||
|
||||
protected abstract CreateTasks createTasks(String[] args, List<ChunkCfg> chunkConfigurations, int variations);
|
||||
|
||||
protected List<String> fileToParagraphs(final Path file, final int chunkNorm, final int chunkVar) {
|
||||
LOGGER.info("Splitting to jsonl " + file + " -> " + chunkNorm + "±" + chunkVar);
|
||||
|
||||
String text;
|
||||
try {
|
||||
text = Utils.normalizeText(Files.readString(file));
|
||||
}
|
||||
catch (final IOException e) {
|
||||
throw new RuntimeException("Cannot read file", e);
|
||||
}
|
||||
|
||||
return Utils.splitParagraphs(text, chunkNorm, chunkVar, false);
|
||||
}
|
||||
|
||||
protected Map<String, String> createCharacterNameMapping(final String paragraph) {
|
||||
final var completion = this.clientService.getClient(ModelCategory.COLD).completion(ModelCategory.COLD,
|
||||
"You are a helpful assistant.",
|
||||
"Extract all names and genders from the following text. Return the result in json format with the following fields: {\"males\" : [ \"Male name 1\", \"Male name 2\", … ], \"females\" : [ \"Female name 1\", \"Female name 2\", … ]}. Don't return a preamble and no explanation:",
|
||||
paragraph);
|
||||
|
||||
final var response = completion.response();
|
||||
final var start = response.indexOf('{');
|
||||
final var end = response.lastIndexOf('}');
|
||||
if (start == -1 || end == -1) {
|
||||
return Map.of();
|
||||
}
|
||||
final var json = response.substring(start, end + 1);
|
||||
|
||||
Names names;
|
||||
final ObjectMapper mapper = new ObjectMapper();
|
||||
try {
|
||||
names = mapper.readValue(json, Names.class);
|
||||
}
|
||||
catch (final JsonProcessingException e) {
|
||||
return Map.of();
|
||||
}
|
||||
|
||||
final Random random = ThreadLocalRandom.current();
|
||||
final Map<String, String> mapping = new HashMap<>();
|
||||
if (names.males != null) {
|
||||
names.males.forEach(name -> mapping.put(name, this.maleNames.get(random.nextInt(this.maleNames.size()))));
|
||||
}
|
||||
if (names.females != null) {
|
||||
names.females.forEach(name -> mapping.put(name, this.femaleNames.get(random.nextInt(this.femaleNames.size()))));
|
||||
}
|
||||
mapping.keySet().removeIf(key -> key.length() < 5);
|
||||
if (mapping.isEmpty()) {
|
||||
return Map.of();
|
||||
}
|
||||
|
||||
LOGGER.debug(mapping.toString());
|
||||
return mapping;
|
||||
}
|
||||
|
||||
protected String renameCharacters(final String paragraph, final Map<String, String> mapping) {
|
||||
var result = paragraph;
|
||||
for (final var entry : mapping.entrySet()) {
|
||||
result = result.replaceAll("(?i)" + entry.getKey(), entry.getValue());
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
protected String paragraphToContextualPrompt(final String paragraph, final String instruction) {
|
||||
final var completion = this.clientService.getClient(ModelCategory.COLD).completion(ModelCategory.COLD,
|
||||
"You are an assistant with great text writing skills.",
|
||||
instruction,
|
||||
paragraph);
|
||||
|
||||
String prompt = completion.response();
|
||||
final int colon = prompt.indexOf(':');
|
||||
if (colon > -1 && colon < prompt.length() + 1) {
|
||||
prompt = prompt.substring(colon + 1);
|
||||
}
|
||||
final int preamble = prompt.indexOf("\n\n");
|
||||
if (preamble > -1 && preamble < prompt.length() + 2) {
|
||||
prompt = prompt.substring(preamble + 2);
|
||||
}
|
||||
prompt = prompt.trim();
|
||||
|
||||
return prompt;
|
||||
}
|
||||
|
||||
protected void writeJsonLine(final ObjectMapper mapper, final OutputStream fos, final Object line) {
|
||||
try {
|
||||
final var jsonb = mapper
|
||||
.writeValueAsString(line)
|
||||
.getBytes(StandardCharsets.UTF_8);
|
||||
fos.write(jsonb);
|
||||
fos.write(NEWLINE);
|
||||
}
|
||||
catch (final IOException e) {
|
||||
LOGGER.error("Cannot convert split to json", e);
|
||||
}
|
||||
}
|
||||
|
||||
protected record ChunkCfg(int chunksNorm, int chunksVar) {
|
||||
}
|
||||
|
||||
protected record Names(List<String> males, List<String> females) {
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user