hf_token="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"

from transformers import AutoModelForCausalLM, AutoTokenizer
# from transformers.src.transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "meta-llama/Llama-2-13b-chat-hf"

# load the model using 4bit quantization (https://huggingface.co/blog/4bit-transformers-bitsandbytes)
model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, use_auth_token=hf_token)
# disable Tensor Parallelism (https://github.com/huggingface/transformers/pull/24906)
model.config.pretraining_tp=1

tokenizer = AutoTokenizer.from_pretrained (model_id, use_auth_token=hf_token)

def extract_lyrics(text):
    start_index = text.find("Verse 1:")
    end_index = text.find("</s>")
    if start_index == -1:
        return text

    text = text[start_index:end_index].strip()

    text = text.replace("</s>", "")

    return text

def generate_lyrics_test(content, sentiment):
    # input_text = "Generate lyrics in standard song format that matches with following requirements. Content should be " + content + ". Genre should be " + genre +  ". Sentiment of the song should be " + sentiment.lower()
    input_text = "Generate lyrics in standard song format that matches with following requirements. Content should be " + content + ". Sentiment of the song should be " + sentiment.lower()

    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
    outputs = model.generate(input_ids, max_length=500)
    return extract_lyrics(tokenizer.decode(outputs[0]))

import gradio as gr

demo = gr.Interface(
    generate_lyrics_test,
    [
        gr.Textbox(lines=5, label = "Content"),
        gr.Dropdown(
            ["Positive", "Negative", "Neutral"], label = "Sentiment"
        ),
    ],
    "text"
)

demo.queue().launch( )

sentiment = 'Positive'
# content = 'create a similar lyric to beat it from michael jackson and write something motivational instead'
content = 'create a similar lyric to beat it from Michael Jackson and write something motivational instead'
# genre = "Rock"
input_text = "Generate lyrics in standard song format that matches with following requirements. Content should be " + content + ". Sentiment of the song should be " + sentiment.lower()
print('input_text:', input_text)