|
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) |
