molbal · July 13, 2024 12:21
diff --git a/text-to-tags.py b/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)
	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)