diff --git a/aiserver.py b/aiserver.py
index 6ece9abb..2ebe148e 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -9433,41 +9433,65 @@ def UI_2_generate_wi(data):
field = data["field"]
existing = data["existing"]
gen_amount = data["genAmount"]
- print(uid, field)
- with open("data/wi_fewshot.txt", "r") as file:
- fewshot_template = file.read()
-
+ # The template to coax what we want from the model
+ extractor_string = ""
+
if field == "title":
- prompt = fewshot_template
for thing in ["type", "desc"]:
if not existing[thing]:
continue
pretty = {"type": "Type", "desc": "Description"}[thing]
- prompt += f"{pretty}: {existing[thing]}\n"
+ extractor_string += f"{pretty}: {existing[thing]}\n"
pretty = "Title"
if existing["desc"]:
# Don't let the model think we're starting a new entry
pretty = "Alternate Title"
- prompt += pretty + ":"
+ extractor_string += pretty + ":"
elif field == "desc":
# MUST be title and type
assert existing["title"]
assert existing["type"]
- prompt = f"{fewshot_template}Title: {existing['title']}\nType: {existing['type']}\nDescription:"
+ extractor_string = f"Title: {existing['title']}\nType: {existing['type']}\nDescription:"
else:
assert False, "What"
+
+ with open("data/wi_fewshot.txt", "r") as file:
+ fewshot_entries = [x.strip() for x in file.read().split("\n\n") if x]
+
+ # Use user's own WI entries in prompt
+ if koboldai_vars.wigen_use_own_wi:
+ fewshot_entries += koboldai_vars.worldinfo_v2.to_wi_fewshot_format(excluding_uid=uid)
- logger.info(prompt)
+ # We must have this amount or less in our context.
+ target = koboldai_vars.max_length - gen_amount - len(tokenizer.encode(extractor_string))
+
+ used = []
+ # Walk the entries backwards until we can't cram anymore in
+ for entry in reversed(fewshot_entries):
+ maybe = [entry] + used
+ maybe_str = "\n\n".join(maybe)
+ possible_encoded = tokenizer.encode(maybe_str)
+ if len(possible_encoded) > target:
+ break
+ yes_str = maybe_str
+ used = maybe
+
+ prompt = f"{yes_str}\n\n{extractor_string}"
+
+ # logger.info(prompt)
+ # TODO: Make single_line mode that stops on newline rather than bans it (for title)
out_text = tpool.execute(
raw_generate,
prompt,
max_new=gen_amount,
single_line=True,
).decoded[0]
- print(f'{out_text}')
+ out_text = utils.trimincompletesentence(out_text.strip())
+
+ socketio.emit("generated_wi", {"uid": uid, "field": field, "out": out_text}, room="UI_2")
@app.route("/generate_raw", methods=["GET"])
def UI_2_generate_raw():
diff --git a/gensettings.py b/gensettings.py
index 1b956827..2e7a2a05 100644
--- a/gensettings.py
+++ b/gensettings.py
@@ -880,6 +880,20 @@ gensettingstf = [
"extra_classes": "simple_ui_only var_sync_alt_user_ui_level",
"ui_level": 0
},
+ {
+ "UI_V2_Only": True,
+ "uitype": "toggle",
+ "unit": "bool",
+ "label": "Native WI Gen",
+ "id": "wigen_use_own_wi",
+ "default": False,
+ "tooltip": "Uses your existing applicable (has title, type, content) World Info entries as inspiration for generated ones.",
+ "menu_path": "World Info",
+ "sub_path": "",
+ "classname": "user",
+ "name": "wigen_use_own_wi",
+ "ui_level": 2
+ },
]
gensettingsik =[{
diff --git a/koboldai_settings.py b/koboldai_settings.py
index 0adf05b7..5c1f1ede 100644
--- a/koboldai_settings.py
+++ b/koboldai_settings.py
@@ -1152,6 +1152,7 @@ class user_settings(settings):
self.img_gen_steps = 30
self.img_gen_cfg_scale = 7.0
self.cluster_requested_models = [] # The models which we allow to generate during cluster mode
+ self.wigen_use_own_wi = False
def __setattr__(self, name, value):
@@ -2471,6 +2472,38 @@ class KoboldWorldInfo(object):
def get_used_wi(self):
return [x['content'] for x in self.world_info if x['used_in_game']]
+
+ def to_wi_fewshot_format(self, excluding_uid: int) -> List[str]:
+ """
+ Returns a list of strings representing applicable (has title, text, and
+ object type) World Info entries. Intended for feeding into the fewshot
+ generator.
+ """
+ the_collection = []
+ for entry in self.world_info.values():
+ if entry["uid"] == excluding_uid:
+ continue
+
+ if not (
+ entry["title"]
+ and entry["manual_text"]
+ and entry["object_type"]
+ ):
+ continue
+
+ processed_desc = entry["manual_text"].replace("\n", " ")
+ while " " in processed_desc:
+ processed_desc = processed_desc.replace(" ", " ")
+ processed_desc = processed_desc.strip()
+
+ the_collection.append(
+ f"Title: {entry['title']}\n" \
+ f"Type: {entry['object_type']}\n"
+ f"Description: {processed_desc}"
+ )
+
+ return the_collection
+
@dataclass
class SavePaths:
diff --git a/static/koboldai.css b/static/koboldai.css
index 5d4f9e9c..c021fde4 100644
--- a/static/koboldai.css
+++ b/static/koboldai.css
@@ -1100,6 +1100,25 @@ td.server_vars {
.world_info_title {
user-select: text;
-moz-user-select:text;
+ margin-bottom: 0px;
+}
+
+.wi-type-leadup {
+ opacity: 0.4;
+}
+
+.world_info_item_type {
+ /* Workaround for Firefox bug: https://stackoverflow.com/a/42283543 */
+ padding: 1px;
+
+ outline: none;
+ border-bottom: 1px solid gray;
+ transition: color 200ms, border-bottom-color 200ms;
+}
+
+.world_info_item_type.bad-input {
+ color: rgb(255, 103, 103);
+ border-bottom-color: red;
}
.world_info_delete {
@@ -1140,6 +1159,24 @@ td.server_vars {
filter: brightness(90%);
}
+.world_info_label_container {
+ display: flex;
+ justify-content: space-between;
+ margin: 1px 5px;
+}
+
+.world_info_label_container > .wi-ui-label {
+ position: relative;
+ top: 2px;
+}
+
+.world_info_label_container > .generate-button {
+ opacity: 0.8;
+ cursor: pointer;
+}
+
+.world_info_label_container > .generate-button:hover { opacity: 1.0; }
+
.tag {
background-color: var(--wi_tag_color);
color: var(--wi_tag_text_color);
diff --git a/static/koboldai.js b/static/koboldai.js
index 2c85c61c..9d4726da 100644
--- a/static/koboldai.js
+++ b/static/koboldai.js
@@ -34,6 +34,7 @@ socket.on("log_message", function(data){process_log_message(data);});
socket.on("debug_message", function(data){console.log(data);});
socket.on("scratchpad_response", recieveScratchpadResponse);
socket.on("show_error_notification", function(data) { reportError(data.title, data.text) });
+socket.on("generated_wi", showGeneratedWIData);
//socket.onAny(function(event_name, data) {console.log({"event": event_name, "class": data.classname, "data": data});});
// Must be done before any elements are made; we track their changes.
@@ -69,6 +70,7 @@ var setup_wi_toggles = [];
var scroll_trigger_element = undefined; //undefined means not currently set. If set to null, it's disabled.
var drag_id = null;
var story_commentary_characters = {};
+var generating_summary = false;
const on_colab = $el("#on_colab").textContent == "true";
// Each entry into this array should be an object that looks like:
@@ -2248,6 +2250,48 @@ function world_info_entry(data) {
document.getElementById("world_info_tags_"+data.uid).classList.remove("hidden");
document.getElementById("world_info_secondtags_"+data.uid).classList.remove("hidden");
}
+
+ const genTypeInput = world_info_card.querySelector(".world_info_item_type");
+ const generateDescButton = world_info_card.querySelector(".wi-lc-text > .generate-button");
+ generateDescButton.addEventListener("click", function() {
+ if (generating_summary) return;
+ let type = genTypeInput.innerText;
+
+ if (!type) {
+ genTypeInput.classList.add("bad-input");
+ return;
+ } else {
+ genTypeInput.classList.remove("bad-input");
+ }
+
+ // TODO: Make type input element
+ generateWIData(data.uid, "desc", title.innerText, type, null, 80);
+ this.innerText = "autorenew";
+ this.classList.add("spinner");
+ });
+
+ genTypeInput.addEventListener("focus", function() {
+ this.classList.remove("bad-input");
+ });
+
+ genTypeInput.addEventListener("keydown", function(event) {
+ if (event.key === "Enter") {
+ event.preventDefault();
+ this.blur();
+ }
+ });
+
+ genTypeInput.addEventListener("blur", function() {
+ this.innerText = this.innerText.trim();
+
+ if (this.innerText == this.getAttribute("old-text")) return;
+ this.setAttribute("old-text", this.innerText);
+
+ world_info_data[data.uid].object_type = this.innerText;
+ send_world_info(data.uid);
+ });
+
+ genTypeInput.innerText = data.object_type;
//$('#world_info_constant_'+data.uid).bootstrapToggle();
//$('#world_info_wpp_toggle_'+data.uid).bootstrapToggle();
@@ -3957,6 +4001,7 @@ function create_new_wi_entry(folder) {
"selective": false,
"wpp": {'name': "", 'type': "", 'format': 'W++', 'attributes': {}},
'use_wpp': false,
+ "object_type": null,
};
var card = world_info_entry(data);
//card.scrollIntoView(false);
@@ -6894,10 +6939,30 @@ for (const proxy of document.querySelectorAll("[sync-proxy-host]")) {
}
function generateWIData(uid, field, title=null, type=null, desc=null, genAmount=80) {
+ if (generating_summary) return;
+ generating_summary = true;
+
socket.emit("generate_wi", {
uid: uid,
field: field,
genAmount: genAmount || 80,
existing: {title: title, type: type, desc: desc}
});
+}
+
+function showGeneratedWIData(data) {
+ generating_summary = false;
+ console.warn(data)
+ const card = $el(`.world_info_card[uid="${data.uid}"]`);
+
+ // Stop spinning!
+ for (const littleRobotFriend of card.querySelectorAll(".generate-button.spinner")) {
+ littleRobotFriend.classList.remove("spinner");
+ littleRobotFriend.innerText = "smart_toy";
+ }
+
+ if (data.field === "desc") {
+ world_info_data[data.uid].manual_text = data.out;
+ send_world_info(data.uid);
+ }
}
\ No newline at end of file
diff --git a/templates/templates.html b/templates/templates.html
index 4a5de1bb..01e2480b 100644
--- a/templates/templates.html
+++ b/templates/templates.html
@@ -10,8 +10,22 @@
>add
![]()
-
+
+
+ is a
+
+
X
@@ -61,7 +75,12 @@
- Text:
+
+ Text
+
+ smart_toy
+
+
@@ -81,6 +100,10 @@
+