diff options
Diffstat (limited to 'describe.py')
| -rw-r--r-- | describe.py | 61 |
1 files changed, 54 insertions, 7 deletions
diff --git a/describe.py b/describe.py index 0e56550..1cce292 100644 --- a/describe.py +++ b/describe.py @@ -34,6 +34,10 @@ PROMPT = ("this is a {kind} from a discord server, its name is \"{hint}\" (the n "the name above can help), what they're doing, and any text written in the image. " "don't guess when to use it, no preamble, no quotes.") +# only for media, the bot sends [media: name] so the file name is what she sees, a hash tells her nothing +NAME_PROMPT = (" answer in exactly two lines. line 1: a short file name for it, 1 to 3 lowercase words joined by underscores, no extension " + "(the character's name if you recognize them, otherwise the main subject). line 2: the description.") + def model_from(script: str) -> str: # same trick as colab_setup.py, read the model name straight out of main.py / gemini.py so they can't drift apart @@ -92,7 +96,7 @@ def video_frame(path: Path, dest: Path) -> bool: def describe(model: str, png: Path, prompt: str) -> str: import ollama r = ollama.chat(model=model, messages=[{"role": "user", "content": prompt, "images": [str(png)]}], options={"temperature": 0.2}) - return clean(r["message"]["content"]) + return strip_think(r["message"]["content"]) def describe_gemini(model: str, png: Path, prompt: str) -> str: @@ -109,18 +113,52 @@ def describe_gemini(model: str, png: Path, prompt: str) -> str: config=types.GenerateContentConfig(safety_settings=safety, temperature=0.2), ) time.sleep(4) # free tier rate limit, rerun the script if some fail, it only does the empty ones - return clean(r.text or "") + return strip_think(r.text or "") + + +def strip_think(text: str) -> str: + return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL) def clean(text: str) -> str: - text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL) - return " ".join(text.split()).strip(' "') + return " ".join(strip_think(text).split()).strip(' "') + + +def split_name(text: str) -> tuple[str, str]: + # media answers come back as "name\ndescription", no name (model ignored the format) = keep the file as it is + lines = [l for l in strip_think(text).splitlines() if l.strip()] + if len(lines) < 2: + return "", clean(text) + name = re.sub(r"[^a-z0-9]+", "_", lines[0].lower()).strip("_")[:40] + return name, clean(" ".join(lines[1:])) -def save_description(kind: str, name: str, text: str, force: bool): +def rename_media(old: str, name: str) -> str: + # keeps the extension, adds _2 _3 if the name is taken. media.txt gets the new name too or the bot forgets the file on restart + suffix = Path(old).suffix + new = f"{name}{suffix}" + n = 2 + while (MEDIA_DIR / new).exists() and new != old: + new = f"{name}_{n}{suffix}" + n += 1 + if new == old: + return old + (MEDIA_DIR / old).rename(MEDIA_DIR / new) + log = MEDIA_DIR / "media.txt" + if log.exists(): + lines = log.read_text(encoding="utf-8").splitlines() + lines = [l[:-len(old)] + new if l.endswith(f"saved media: {old}") else l for l in lines] + log.write_text("\n".join(lines) + "\n", encoding="utf-8") + return new + + +def save_description(kind: str, name: str, text: str, force: bool, new_name: str = None): # read it fresh every time, the bot adds new entries to this file while we run data = json.loads(REACTIONS_FILE.read_text(encoding="utf-8")) if name in data.get(kind, {}) and (force or not data[kind][name]): + if new_name and new_name != name: + data[kind][new_name] = data[kind].pop(name) + name = new_name data[kind][name] = text REACTIONS_FILE.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") @@ -161,14 +199,23 @@ def main(): try: hint = re.sub(r"^\d+_", "", Path(name).stem) if kind == "media" else name # media is saved as id_filename prompt = PROMPT.format(kind=kind.rstrip("s") if kind != "media" else "media file", hint=hint) + if kind == "media": + prompt += NAME_PROMPT text = (describe_gemini if args.gemini else describe)(model, png, prompt) except Exception as e: print(f"skip {kind}/{name}: describe failed: {e}") continue + new_name = None + if kind == "media": + new_name, text = split_name(text) + else: + text = clean(text) if not text: continue - save_description(kind, name, text, args.all) - print(f"{kind}/{name}: {text}") + if new_name: + new_name = rename_media(name, new_name) + save_description(kind, name, text, args.all, new_name) + print(f"{kind}/{name}{f' -> {new_name}' if new_name and new_name != name else ''}: {text}") if __name__ == "__main__": |