Skip to content

Commit

Permalink
program loops, better prompt override
Browse files Browse the repository at this point in the history
  • Loading branch information
@HAV0X1014
HAV0X1014 committed Sep 17, 2024
1 parent 7d0e128 commit 222c10b
Show file tree
Hide file tree
Showing 3 changed files with 128 additions and 102 deletions.
32 changes: 28 additions & 4 deletions src/main/java/GeminiAI.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

public class GeminiAI {
private static final MediaType JSON = MediaType.get("application/json; charset=utf-8");
public static String send(String dialog, String untranslatedNames, String translatedNames) {
public static String send(String dialog, String mainCharacter) {
OkHttpClient client = new OkHttpClient.Builder().connectTimeout(60, TimeUnit.SECONDS).readTimeout(120, TimeUnit.SECONDS).build();
JSONObject payload = new JSONObject(); //holds the whole object of JSON sent to google
JSONArray contents = new JSONArray(); //holds the text content sent to the AI
Expand All @@ -23,7 +23,7 @@ public static String send(String dialog, String untranslatedNames, String transl
"- Be specific about the tone of the characters. Describe their expressions and tone in detail.\n" +
"- This story happens in the main story of Kemono Friends 3, where human girls with animal features go on adventures, and fight against Celliens.\n";
} else {
prompt = ConfigHandler.getString("PromptOverride");
prompt = ConfigHandler.getString("PromptOverride").replaceAll("\\{char}", mainCharacter);
}
/*
//this would have been for a "system" prompt but apparently gemini doesnt have that.
Expand All @@ -50,14 +50,38 @@ public static String send(String dialog, String untranslatedNames, String transl
//this hunk of shit puts the text into googles special little array with an object inside or something
//i hate google's stupid fucking API formatting

JSONArray safetySettings = new JSONArray();
JSONObject catHar = new JSONObject();
catHar.put("category", "HARM_CATEGORY_HARASSMENT");
catHar.put("threshold", "BLOCK_NONE");

JSONObject catHtSp = new JSONObject();
catHtSp.put("category", "HARM_CATEGORY_HATE_SPEECH");
catHtSp.put("threshold", "BLOCK_NONE");

JSONObject catSxExp = new JSONObject();
catSxExp.put("category", "HARM_CATEGORY_SEXUALLY_EXPLICIT");
catSxExp.put("threshold", "BLOCK_NONE");

JSONObject catDngCnt = new JSONObject();
catDngCnt.put("category", "HARM_CATEGORY_DANGEROUS_CONTENT");
catDngCnt.put("threshold", "BLOCK_NONE");

safetySettings.put(catHar);
safetySettings.put(catHtSp);
safetySettings.put(catSxExp);
safetySettings.put(catDngCnt);

payload.put("safetySettings", safetySettings);
payload.put("contents",contents);

RequestBody requestBody = RequestBody.create(JSON, payload.toString());
String output;
String responseContent = "";

try {
//to change model, change the gemini-1.5-pro-latest thing to whichever model you want.
URL url = new URL("https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro-latest:generateContent?key=" + ConfigHandler.getString("GeminiAPIKey"));
//to change model, change the gemini-1.5-flash-latest thing to whichever model you want.
URL url = new URL("https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent?key=" + ConfigHandler.getString("GeminiAPIKey"));
Request request = new Request.Builder().url(url).post(requestBody).build();
try (Response resp = client.newCall(request).execute()) {
responseContent = resp.body().string();
Expand Down
14 changes: 3 additions & 11 deletions src/main/java/Summarize.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,17 +38,9 @@ public static void summarize(File[] files, String folderPath, String id, String
StringBuilder names = new StringBuilder();
List<Object> charaNamesList = charaNames.toList().stream().distinct().collect(Collectors.toList());

JSONArray jsonArray = new JSONArray(FileHandler.read("CHARA_DATA.json"));
// Populate the map with data from the chara data list
Map<String, JSONObject> nameMap = new HashMap<>();
for (int i = 0; i < jsonArray.length(); i++) {
JSONObject jsonObject = jsonArray.getJSONObject(i);
String name = jsonObject.getString("name");
nameMap.put(name, jsonObject);
}
//look up the name of the character in the Map and match that to their respective friend ID and english name
for (Object name : charaNamesList.toArray()) {
JSONObject result = nameMap.get(name.toString());
JSONObject result = SummaryMain.nameMap.get(name.toString());
if (result != null) {
String nameEn = result.getString("nameEn");
names.append(name).append("|").append(nameEn).append(", ");
Expand All @@ -73,7 +65,7 @@ public static void summarize(File[] files, String folderPath, String id, String
String cleanedSentence = sentence.toString().replaceAll("<.*?>", "");
//if the name we are looking up (the current speaking character's name) ISNT in
//the name map, then dont add it to the dialog pair. (cellien, human characters, etc)
JSONObject result = nameMap.get(charName);
JSONObject result = SummaryMain.nameMap.get(charName);
if (result == null) {
dialog = charName + ": " + cleanedSentence + "\n";
} else {
Expand All @@ -92,7 +84,7 @@ public static void summarize(File[] files, String folderPath, String id, String
System.out.println("---");
System.out.println("Summarizing " + id + "...");
//String translatedNames = Translate.translator(names); //translate the names with google translate first UNUSED
String summarizedScene = GeminiAI.send(fullScene.toString(), names.toString(), ""); //send the scenarios to the AI
String summarizedScene = GeminiAI.send(fullScene.toString(), SummaryMain.idMap.get(Integer.parseInt(id)).getString("nameEn")); //send the scenarios to the AI
System.out.println(summarizedScene);

//now write the summary into a file
Expand Down
184 changes: 97 additions & 87 deletions src/main/java/SummaryMain.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
public class SummaryMain {
public static JSONObject config;
public static final File SummaryDirs = new File("summary/untranslated/");
public static Map<Integer, JSONObject> idMap = new HashMap<>();
public static Map<String, JSONObject> nameMap = new HashMap<>();
public static JSONArray charaData = new JSONArray(FileHandler.read("CHARA_DATA.json"));
public static void main(String[] args) {
Scanner sc = new Scanner(System.in);
String folderPath;
Expand All @@ -19,102 +22,109 @@ public static void main(String[] args) {
SummaryDirs.mkdirs();
}

//user input and validation
System.out.println("Choose which story you want to summarize. Type \"main\", \"main2\", \"main3\", \"main4\", \"another\", \"event\" or a Friend ID.\n" +
"(To summarize a Friend story, input the ID of the Friend, format \"3\" or \"179\". \"event\" requires an ID for the event.)");
String choice = sc.nextLine();
while (!choice.equals("main") && !choice.equals("main2") && !choice.equals("main3") && !choice.equals("main4") && !choice.equals("another") && !choice.equals("event") && !isInt(choice)) {
System.out.println("Incorrect selection. Choose one of the above choices.");
choice = sc.nextLine();
// Populate the map with data from the chara data list
for (int i = 0; i < charaData.length(); i++) {
JSONObject jsonObject = charaData.getJSONObject(i);
Integer friendID = jsonObject.getInt("id");
idMap.put(friendID, jsonObject);
}
// Populate the map with data from the chara data list
for (int i = 0; i < charaData.length(); i++) {
JSONObject jsonObject = charaData.getJSONObject(i);
String name = jsonObject.getString("name");
nameMap.put(name, jsonObject);
}
while (true) {
//user input and validation
System.out.println("Choose which story you want to summarize. Type \"main\", \"main2\", \"main3\", \"main4\", \"another\", \"event\" or a Friend ID.\n" +
"(To summarize a Friend story, input the ID of the Friend, format \"3\" or \"179\". \"event\" requires an ID for the event.)\n" +
"[Use ctrl + c to quit]");
String choice = sc.nextLine();
while (!choice.equals("main") && !choice.equals("main2") && !choice.equals("main3") && !choice.equals("main4") && !choice.equals("another") && !choice.equals("event") && !isInt(choice)) {
System.out.println("Incorrect selection. Choose one of the above choices.");
choice = sc.nextLine();
}

//read the config here, as the user should have entered it by now.
String configString = FileHandler.read("config.json");
config = new JSONObject(configString.trim());
//read the config here, as the user should have entered it by now.
String configString = FileHandler.read("config.json");
config = new JSONObject(configString.trim());

switch (choice) {
case "main":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 10; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main_");
}
break;
case "main2":
folderPath = "main/";
for (int loopOverEveryChapter = 1; loopOverEveryChapter < 6; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_002_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main2_");
}
break;
case "main3":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 8; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_003_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main3_");
}
break;
case "main4":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 5; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_004_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main4_");
}
break;
case "another":
folderPath = "another/";
for (int loopOverEveryChapter = 1; loopOverEveryChapter < 7; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_a_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "another_");
}
break;
case "event": {
folderPath = "event/";
System.out.println("Input the ID of the event you want to summarize.");
String id = sc.nextLine();

switch (choice) {
case "main":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 10; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main_");
}
break;
case "main2":
folderPath = "main/";
for (int loopOverEveryChapter = 1; loopOverEveryChapter < 6; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_002_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main2_");
}
break;
case "main3":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 8; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_003_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main3_");
}
break;
case "main4":
folderPath = "main/";
for (int loopOverEveryChapter = 0; loopOverEveryChapter < 5; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_m_004_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "main4_");
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_e_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "event_");
break;
}
break;
case "another":
folderPath = "another/";
for (int loopOverEveryChapter = 1; loopOverEveryChapter < 7; loopOverEveryChapter++) {
String id = String.format("%02d", loopOverEveryChapter);
default: {
folderPath = "char/";
int idUnformattedInt = Integer.parseInt(choice);
String id = String.format("%04d", idUnformattedInt);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_a_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "another_");
}
break;
case "event": {
folderPath = "event/";
System.out.println("Input the ID of the event you want to summarize.");
String id = sc.nextLine();

File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_e_" + id + "_") && name.endsWith(".prefab.json"));
Summarize.summarize(files, folderPath, id, "event_");
break;
}
default: {
folderPath = "char/";
int idUnformattedInt = Integer.parseInt(choice);
String id = String.format("%04d", idUnformattedInt);
File folder = new File(folderPath);
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_c_" + id + "_") && name.endsWith(".prefab.json"));
File[] files = folder.listFiles((dir, name) -> name.startsWith("scenario_c_" + id + "_") && name.endsWith(".prefab.json"));

JSONArray jsonArray = new JSONArray(FileHandler.read("CHARA_DATA.json"));
// Populate the map with data from the chara data list
Map<Integer, JSONObject> idMap = new HashMap<>();
for (int i = 0; i < jsonArray.length(); i++) {
JSONObject jsonObject = jsonArray.getJSONObject(i);
Integer friendID = jsonObject.getInt("id");
idMap.put(friendID, jsonObject);
}
String nameEn = "NO-MATCH";
//look up the id of the character in the Map and match that to their respective name and english name
for (Object friend : idMap.keySet().toArray()) {
JSONObject result = idMap.get(Integer.parseInt(friend.toString()));
if (result.getInt("id") == Integer.parseInt(id)) {
int friendID = result.getInt("id");
nameEn = result.getString("nameEn");
System.out.println("ID: " + friendID + ", NameEN: " + nameEn);
String nameEn = "NO-MATCH";
//look up the id of the character in the Map and match that to their respective name and english name
for (Object friend : idMap.keySet().toArray()) {
JSONObject result = idMap.get(Integer.parseInt(friend.toString()));
if (result.getInt("id") == Integer.parseInt(id)) {
int friendID = result.getInt("id");
nameEn = result.getString("nameEn");
System.out.println("ID: " + friendID + ", NameEN: " + nameEn);
}
}
Summarize.summarize(files, folderPath, id, nameEn + "_char_");
break;
}
Summarize.summarize(files, folderPath, id, nameEn + "_char_");
break;
}
}
}
Expand Down

0 comments on commit 222c10b

Please sign in to comment.