Add the app

app.py

@@ -0,0 +1,63 @@
+import gradio as gr
+from functions import CAT_ARTIST, CAT_CHARACTER, CAT_COPYRIGHT, CAT_GENERAL, CAT_LORE, CAT_META, CAT_SPECIES, PromptBuilder, parse_tag, parse_tags, related_tags
+
+def query_tag(tag: str, category: int):
+    if category == -1:
+        category = None 
+    return related_tags(parse_tag(tag), category=category)
+
+def generate_prompt(include: str, focus: str, exclude: str, avoid: str, skip: str, rating: str, general: int, artist: int, species: int, copyright: int, character: int, meta: int) -> str:
+    try:
+        builder = PromptBuilder(skip=list(parse_tags(skip)), min_posts=50, rating=rating)
+        for tag in parse_tags(include):
+            builder = builder.include(tag)
+        for tag in parse_tags(focus):
+            builder = builder.focus(tag)
+        for tag in parse_tags(exclude):
+            builder = builder.exclude(tag)
+        for tag in parse_tags(avoid):
+            builder = builder.avoid(tag)
+        if artist > 0:
+            builder = builder.pick(CAT_ARTIST, artist, 10)
+        if species > 0:
+            builder = builder.pick(CAT_SPECIES, species, 10)
+        if copyright > 0:
+            builder = builder.pick(CAT_COPYRIGHT, copyright, 10)
+        if character > 0:
+            builder = builder.pick(CAT_CHARACTER, character, 10)
+        if meta > 0:
+            builder = builder.pick(CAT_META, meta, 10)
+        if general > 0:
+            builder = builder.pick(CAT_GENERAL, general, 50)
+        return builder.get_one()
+    except Exception as e:
+        return str(e)
+
+with gr.Blocks() as demo:
+    with gr.Tab("Tag Explorer"):
+        tag = gr.Textbox(label="Tag")
+        category = gr.Dropdown(label="Category", choices=[("All", -1), ("General", CAT_GENERAL), ("Artist", CAT_ARTIST), ("Copyright", CAT_COPYRIGHT), ("Character", CAT_CHARACTER), ("Species", CAT_SPECIES), ("Meta", CAT_META), ("Lore", CAT_LORE)], value=-1)
+        query = gr.Button("Query")
+        output = gr.Dataframe()
+        query.click(fn=query_tag, inputs=[tag, category], outputs=output)
+    with gr.Tab("Prompt Expander"):
+        include = gr.Textbox(label="Positive Prompt - Start with these tags and weight picks toward them")
+        focus = gr.Textbox(label="Focus - Used for picks but not necessarily added to the prompt")
+        exclude = gr.Textbox(label="Negative Prompt - Put these in the negative prompt and weight picks against them")
+        avoid = gr.Textbox(label="Other Negatives - Weighted against picks but not put in the negatives")
+        skip = gr.Textbox(label="Skip - Never pick these tags")
+        rating = gr.Dropdown(label="Rating Limit (not 100% reliable)", choices=[("Safe", "s"), ("Questionable", "q"), ("Explicit", "e")], value="s")
+        with gr.Accordion(label="Tag Counts"):
+            with gr.Row():
+                general = gr.Number(5, label="General", precision=0, minimum=0, maximum=20)
+                artist = gr.Number(0, label="Artist", precision=0, minimum=0, maximum=5)
+                species = gr.Number(0, label="Species", precision=0, minimum=0, maximum=5)
+                copyright = gr.Number(0, label="Copyright", precision=0, minimum=0, maximum=5)
+                character = gr.Number(0, label="Character", precision=0, minimum=0, maximum=5)
+                meta = gr.Number(0, label="Meta", precision=0, minimum=0, maximum=5)
+
+        generate = gr.Button("Generate")
+        output = gr.Textbox(label="Prompt")
+        generate.click(fn=generate_prompt, inputs=[include, focus, exclude, avoid, skip, rating, general, artist, species, copyright, character, meta], outputs=output)
+
+demo.launch()

functions.py

@@ -0,0 +1,267 @@
+import pandas
+import numpy
+import pandas.io.formats.style
+import random
+import functools
+from typing import Callable, Literal
+
+DATA_FOLDER = "."
+
+CAT_GENERAL = 0
+CAT_ARTIST = 1
+CAT_UNUSED = 2
+CAT_COPYRIGHT = 3
+CAT_CHARACTER = 4
+CAT_SPECIES = 5
+CAT_INVALID = 6
+CAT_META = 7
+CAT_LORE = 8
+
+CATEGORY_COLORS = {
+    CAT_GENERAL: "#b4c7d9",
+    CAT_ARTIST: "#f2ac08",
+    CAT_UNUSED: "#ff3d3d",
+    CAT_COPYRIGHT: "#d0d",
+    CAT_CHARACTER: "#0a0",
+    CAT_SPECIES: "#ed5d1f",
+    CAT_INVALID: "#ff3d3d",
+    CAT_META: "#fff",
+    CAT_LORE: "#282"
+}
+
+def get_feather(filename: str) -> pandas.DataFrame:
+    return pandas.read_feather(f"{DATA_FOLDER}/{filename}.feather")
+
+tags = get_feather("tags")
+posts_by_tag = get_feather("posts_by_tag").set_index("tag_id")
+tags_by_post = get_feather("tags_by_post").set_index("post_id")
+tag_ratings = get_feather("tag_ratings")
+implications = get_feather("implications")
+tags_by_name = tags.copy(deep=True)
+tags_by_name.set_index("name", inplace=True)
+tags.set_index("tag_id", inplace=True)
+
+@functools.cache
+def get_related_tags(targets: tuple[str, ...], exclude: tuple[str, ...] = (), samples: int = 100_000) -> pandas.DataFrame:
+    these_tags = tags_by_name.loc[list(targets)]
+    posts_with_these_tags = posts_by_tag.loc[these_tags["tag_id"]].map(set).groupby(lambda x: True).agg(lambda x: set.intersection(*x))["post_id"][True]
+    if (len(exclude) > 0):
+        excluded_tags = tags_by_name.loc[list(exclude)]
+        posts_with_excluded_tags = posts_by_tag.loc[excluded_tags["tag_id"]].map(set).groupby(lambda x: True).agg(lambda x: set.union(*x))["post_id"][True]
+        posts_with_these_tags = posts_with_these_tags - posts_with_excluded_tags
+    total_post_count_together = len(posts_with_these_tags)
+    sample_posts = random.sample(list(posts_with_these_tags), samples) if total_post_count_together > samples else list(posts_with_these_tags)
+    post_count_together = len(sample_posts)
+    sample_ratio = post_count_together / total_post_count_together
+    tags_in_these_posts = tags_by_post.loc[sample_posts]
+    counts_in_these_posts = tags_in_these_posts["tag_id"].explode().value_counts().rename("overlap")
+    summaries = pandas.DataFrame(counts_in_these_posts).join(tags[tags["post_count"]>0], how="right").fillna(0)
+    summaries["overlap"] = numpy.minimum(summaries["overlap"] / sample_ratio, summaries["post_count"])
+    summaries = summaries[["category", "name", "overlap", "post_count"]]
+    # Old "interestingness" value, didn't give as good results as an actual statistical technique, go figure. Code kept for curiosity's sake.
+    #summaries["interestingness"] = summaries["overlap"].pow(2) / (total_post_count_together * summaries["post_count"])
+    # Phi coefficient stuff.
+    n = float(len(tags_by_post))
+    n11 = summaries["overlap"]
+    n1x = float(total_post_count_together)
+    nx1 = summaries["post_count"].astype("float64")
+    summaries["correlation"] = (n * n11 - n1x * nx1) / numpy.sqrt(n1x * nx1 * (n - n1x) * (n - nx1))
+    return summaries
+
+def format_tags(styler: pandas.io.formats.style.Styler):
+    styler.apply(lambda row: numpy.where(row.index == "name", "color:"+CATEGORY_COLORS[row["category"]], ""), axis=1)
+    styler.hide(level=0)
+    styler.hide("category",axis=1)
+    if 'overlap' in styler.data:
+        styler.format("{:.0f}".format, subset=["overlap"])
+    if 'correlation' in styler.data:
+        styler.format("{:.2f}".format, subset=["correlation"])
+        styler.background_gradient(vmin=-1.0, vmax=1.0, cmap="RdYlGn", subset=["correlation"])
+    if 'score' in styler.data:
+        styler.format("{:.2f}".format, subset=["score"])
+        styler.background_gradient(vmin=-1.0, vmax=1.0, cmap="RdYlGn", subset=["score"])
+    return styler
+
+def related_tags(*targets: str, exclude: tuple[str, ...] = (), category: int = None, samples: int = 100_000, min_overlap: int = 5, min_posts: int = 20, top: int = 30, bottom: int = 0) -> pandas.DataFrame:
+    result = get_related_tags(targets, exclude=exclude, samples=samples)
+    if category != None:
+        result = result[result["category"] == category]
+    result = result[~result["name"].isin(targets)]
+    result = result[result["overlap"] >= min_overlap]
+    result = result[result["post_count"] >= min_posts]
+    top_part = result.sort_values("correlation", ascending=False)[:top]
+    bottom_part = result.sort_values("correlation", ascending=True)[:bottom].sort_values("correlation", ascending=False)
+    return pandas.concat([top_part, bottom_part]).style.pipe(format_tags)
+
+def implications_for(*subjects: str, seen: set[str] = None):
+    if seen is None:
+        seen = set()
+    for subject in subjects:
+        found = tags.loc[list(implications[implications["antecedent_id"] == tags_by_name.loc[subject, "tag_id"]].loc[:,"consequent_id"]), "name"].values
+        for f in found:
+            if f in seen:
+                pass
+            else:
+                yield f
+                seen.add(f)
+                yield from implications_for(f, seen=seen)
+
+def parse_tag(potential_tag: str):
+    potential_tag = potential_tag.strip().replace(" ", "_").replace("\\(", "(").replace("\\)", ")")
+    if potential_tag == "":
+        return None
+    elif potential_tag in tags_by_name.index:
+        return potential_tag
+    elif potential_tag.startswith("by_") and potential_tag[3:] in tags_by_name.index:
+        return potential_tag[3:]
+    else:
+        print(f"Couldn't find tag '{potential_tag}', skipping it.")
+
+def parse_tags(*parts: str):
+    for part in parts:
+        for potential_tag in part.split(","):
+            tag = parse_tag(potential_tag)
+            if tag is not None:
+                yield tag
+
+def add_suggestions(suggestions: pandas.DataFrame, new_tags: str | list[str], multiplier: int, samples : int, min_posts: int, rating: Literal['s', 'q', 'e']):
+    if isinstance(new_tags, str):
+        new_tags = [new_tags]
+    for new_tag in new_tags:
+        related = get_related_tags((new_tag,), samples=samples)
+        # Implementing the rating filter this way is horribly inefficient, fix it later
+        if rating == 's':
+            related = related.join(tag_ratings.set_index("tag_id"), on="tag_id")
+            related["post_count"] = related["s"]
+            related = related.drop("s", axis=1)
+            related = related.drop("q", axis=1)
+            related = related.drop("e", axis=1)
+        elif rating == 'q':
+            related = related.join(tag_ratings.set_index("tag_id"), on="tag_id")
+            related["post_count"] = related["s"] + related["q"]
+            related = related.drop("s", axis=1)
+            related = related.drop("q", axis=1)
+            related = related.drop("e", axis=1)
+        related = related[related["post_count"] >= min_posts]
+        if suggestions is None:
+            suggestions = related.rename(columns={"correlation": "score"})
+        else:
+            suggestions = suggestions.join(related, rsuffix="r")
+            # This is a totally made up way to combine correlations. It keeps them from going outside the +/- 1 range, which is nice. It also makes older
+            # tags less important every time newer ones are added. That could be considered a feature or not.
+            suggestions["score"] = numpy.real(numpy.power((numpy.sqrt(suggestions["score"] + 0j) + numpy.sqrt(multiplier * suggestions["correlation"] + 0j)) / 2, 2))
+    return suggestions[["category", "name", "post_count", "score"]]
+
+
+
+def pick_tags(suggestions: pandas.DataFrame, category: int, count: int, from_top: int, excluding: list[str], weighted: bool = True):
+    options = suggestions[(True if category is None else suggestions["category"] == category) & (suggestions["score"] > 0) & ~suggestions["name"].isin(excluding)].sort_values("score", ascending=False)[:from_top]
+    if weighted:
+        values = list(options["name"].values)
+        weights = list(options["score"].values)
+        choices = []
+        for _ in range(count):
+            choice = random.choices(population=values, weights=weights, k=1)[0]
+            weights.pop(values.index(choice))
+            values.remove(choice)
+            choices.append(choice)
+        return choices
+    else:
+        return random.sample(list(options["name"].values), count)
+
+def tag_to_prompt(tag: str) -> str:
+    if (tags_by_name.loc[tag]["category"] == CAT_ARTIST):
+        tag = "by " + tag
+    return tag.replace("_", " ").replace("(" , "\\(").replace(")" , "\\)")
+
+# A lambda in a for loop doesn't capture variables the way I want it to, so this is a method now
+def add_suggestions_later(suggestions: pandas.DataFrame, new_tags: str | list[str], multiplier: int, samples: int, min_posts: int, rating: Literal['s', 'q', 'e']):
+    return lambda: add_suggestions(suggestions, new_tags, multiplier, samples, min_posts, rating)
+
+
+Prompt = tuple[list[str], list[str], Callable[[], pandas.DataFrame]]
+
+class PromptBuilder:
+    prompts: list[Prompt]
+    samples: int
+    min_posts: int
+    rating: Literal['s', 'q', 'e']
+    skip_list: list[str]
+
+    def __init__(self, prompts = [([],[],lambda: None)], skip=[], samples = 100_000, min_posts = 20, rating: Literal['s', 'q', 'e'] = 'e'):
+        self.prompts = prompts
+        self.samples = samples
+        self.min_posts = min_posts
+        self.rating = rating
+        self.skip_list = skip
+
+    def include(self, tag: str):
+        return PromptBuilder(prompts=[
+            (tag_list + [tag], negative_list, add_suggestions_later(suggestions(), tag, 1, self.samples, self.min_posts, self.rating))
+            for (tag_list, negative_list, suggestions) in self.prompts
+        ], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def focus(self, tag: str):
+        return PromptBuilder(prompts=[
+            (tag_list, negative_list, add_suggestions_later(suggestions(), tag, 1, self.samples, self.min_posts, self.rating))
+            for (tag_list, negative_list, suggestions) in self.prompts
+        ], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def exclude(self, tag: str):
+        return PromptBuilder(prompts=[
+            (tag_list, negative_list + [tag], add_suggestions_later(suggestions(), tag, -1, self.samples, self.min_posts, self.rating))
+            for (tag_list, negative_list, suggestions) in self.prompts
+        ], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def avoid(self, tag: str):
+        return PromptBuilder(prompts=[
+            (tag_list, negative_list, add_suggestions_later(suggestions(), tag, -1, self.samples, self.min_posts, self.rating))
+            for (tag_list, negative_list, suggestions) in self.prompts
+        ], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def pick(self, category: int, count: int, from_top: int):
+        new_prompts = self.prompts
+        for _ in range(count):
+            new_prompts = [
+                (tag_list + [tag], negative_list, add_suggestions_later(s, tag, 1, self.samples, self.min_posts, self.rating))
+                for (tag_list, negative_list, suggestions) in new_prompts
+                for s in (suggestions(),)
+                for tag in pick_tags(s, category, 1, from_top, tag_list + negative_list + self.skip_list)
+            ]
+        return PromptBuilder(new_prompts, samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def foreach_pick(self, category: int, count: int, from_top: int):
+        return PromptBuilder(prompts=[
+            (tag_list + [tag], negative_list, add_suggestions_later(s, tag, 1, self.samples, self.min_posts, self.rating))
+            for (tag_list, negative_list, suggestions) in self.prompts
+            for s in (suggestions(),)
+            for tag in pick_tags(s, category, count, from_top, tag_list + negative_list + self.skip_list)
+        ], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+    
+    def pick_fast(self, category: int, count: int, from_top: int):
+        prompts = []
+        for (tag_list, negative_list, suggestions) in self.prompts:
+            s = suggestions()
+            new_tags = pick_tags(s, category, count, from_top, tag_list + negative_list + self.skip_list)
+            prompts.append((tag_list + new_tags, negative_list, add_suggestions_later(s, new_tags, 1, self.samples, self.min_posts, self.rating)))
+        return PromptBuilder(prompts=prompts, samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def branch(self, count: int):
+        return PromptBuilder(prompts=[prompt for prompt in self.prompts for _ in range(count)], samples=self.samples, min_posts=self.min_posts, skip=self.skip_list, rating=self.rating)
+
+    def build(self):
+        for (tag_list, negative_list, _) in self.prompts:
+            positive_prompt = ", ".join([ tag_to_prompt(tag) for tag in tag_list])
+            negative_prompt = ", ".join([ tag_to_prompt(tag) for tag in negative_list])
+            if negative_prompt:
+                yield f"{positive_prompt}\nNegative prompt: {negative_prompt}"
+            else:
+                yield positive_prompt
+
+    def print(self):
+        for prompt in self.build():
+            print(prompt)
+
+    def get_one(self):
+        for prompt in self.build():
+            return prompt