import replicate

def describe_noun_llama_3_70b(noun: str) -> str:
    response = ""
    # The meta/meta-llama-3-8b model can stream output as it's running.
    for event in replicate.stream(
            "meta/meta-llama-3-70b-instruct",
            input={
                "top_k": 0,
                "top_p": 0.95,
                "prompt": f"You will describe an art piece based on this noun {noun}.\nThe description should be suitable for creating a cartoon-like, emoji-style, or minimalist SVG design that functions effectively as an icon.\nSince the icon will be small, focus on bold, simple details that make the image visually distinctive and easy to recognize.\nEmphasize strong shapes, clear contrasts, and vibrant colors to ensure the design is memorable and stands out at a small scale.\nYour response will feed directly into a text-to-image model like DALL·E 3, so provide precise and concrete visual descriptions.\nAvoid unnecessary complexity or fine details that may be lost when scaled down.\nIf the word provided is not a noun or is inappropriate, respond with 'invalid'.\nYour only output should be a description or 'invalid'. DO NOT ADD ANYTHING EXTRA.",
                "max_tokens": 512,
                "temperature": 0.7,
                "system_prompt": "You are a helpful assistant",
                "length_penalty": 1,
                "max_new_tokens": 512,
                "stop_sequences": "<|end_of_text|>,<|eot_id|>",
                "prompt_template": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
                "presence_penalty": 0,
                "log_performance_metrics": False
            },
    ):
        response += str(event)

    return response

def describe_noun_llama_3_8b(noun: str) -> str:
    response = ""
    # The meta/meta-llama-3-8b model can stream output as it's running.
    for event in replicate.stream(
            "meta/meta-llama-3-8b-instruct",
            input={
                "top_k": 0,
                "top_p": 0.95,
                "prompt": f"You will describe an art piece based on this noun {noun}.\nThe description should be suitable for creating a cartoon-like, emoji-style, or minimalist SVG design that functions effectively as an icon.\nSince the icon will be small, focus on bold, simple details that make the image visually distinctive and easy to recognize.\nEmphasize strong shapes, clear contrasts, and vibrant colors to ensure the design is memorable and stands out at a small scale.\nYour response will feed directly into a text-to-image model like DALL·E 3, so provide precise and concrete visual descriptions.\nAvoid unnecessary complexity or fine details that may be lost when scaled down.\nIf the word provided is not a noun or is inappropriate, respond with 'invalid'.\nYour only output should be a description or 'invalid'. DO NOT ADD ANYTHING EXTRA.",
                "max_tokens": 512,
                "temperature": 0.7,
                "system_prompt": "You are a helpful assistant",
                "length_penalty": 1,
                "max_new_tokens": 512,
                "stop_sequences": "<|end_of_text|>,<|eot_id|>",
                "prompt_template": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
                "presence_penalty": 0,
                "log_performance_metrics": False
            },
    ):
       response += str(event)

    return response

def deduplicate(nouns_list: list[str]) -> list[str]:
    return list(set([noun.strip() for noun in nouns_list]))

if __name__ == "__main__":
    #completion = describe_noun("blank")
    with open('../data/imageable_gptfilter.txt', 'r') as fp:
        nouns = fp.read()

    with open('../data/imageable_deduplicate.txt', 'w') as fp:
        fp.write("\n".join(deduplicate(nouns.split("\n"))))

    #completion = add_nouns(noun_list=nouns)

    #output_file = '../data/imageable_gpt_addition003.txt'
    #with open(output_file, 'w') as fp:
    #    fp.write(completion.content)
    #    print(f"written to {output_file}")

