functions.py

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: "#808080", 
 CAT_ARTIST: "#f2ac08", 
 CAT_UNUSED: "#ff3d3d", 
 CAT_COPYRIGHT: "#d0d", 
 CAT_CHARACTER: "#0a0", 
 CAT_SPECIES: "#ed5d1f", 
 CAT_INVALID: "#ff3d3d", 
 CAT_META: "#04f", 
 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