text-to-tags.py

import os
from tqdm import tqdm
import ollama

directory = 'C:\\tools\\training\\illustration_pending'
model_name = 'llama3:8b-instruct-q6_K'


def rewrite_description(description):
    """
    Use the local LLM to rewrite the description into tag-based format.
    """
    few_shot_examples = (
        "Example 1:\n"
        "Input: In this image we can see a painting of a hen and a ball on the mat. Thóe background is dark.\n"
        "Output: hen, mat, egg\n\n"
        "Example 2:\n"
        "Input: In this painting we can see a painting of two persons. The background is dark.\n"
        "Output: 2people, dark_background\n\n"
        "Example 3:\n"
        "Input: This is an animated image. In this image we can see some insects and grass on the ground.\n"
        "Output: insect, grass, ground\n\n"
    )

    prompt = (
        f"{few_shot_examples}"
        f"Rewrite the following description into a tag-based format. Use commas to separate tags. Remove mediums such as painting, cartoon:\n"
        f"Input: {description}\n"
        f"Output:"
    )

    response = ollama.chat(model=model_name, messages=[{'role': 'user', 'content': prompt}])
    return response['message']['content'].strip()


# Iterate over all txt files in the directory
for filename in tqdm(os.listdir(directory)):
    if filename.endswith('.txt'):
        file_path = os.path.join(directory, filename)

        # Read the content of the file
        with open(file_path, 'r') as file:
            original_content = file.read()
            file.close()

        # Rewrite the content using the local LLM
        new_content = rewrite_description(original_content)

        # Print the 'from' and 'to' values
        print(f"\nFile: {file_path}")
        print(f"\nFrom: {original_content}")
        print(f"To: {new_content}\n")

        # Clear the original content and save the new content
        with open(file_path, 'w') as file:
            file.write(new_content)