from ArtistCoherencyModel import ArtistCoherencyModel
import streamlit as st
import pandas as pd

from LyricGeneratorModel import LyricGeneratorModel

artists_df = pd.read_csv("artists.csv")
artist_names_list = list(artists_df["name"])


artist_name_input = st.selectbox("Artist", artist_names_list)
song_title = st.text_input("Song Title")
song_description = st.text_area("Song Description")
submit_button = st.button("Submit")

lyric_evaluator_model = ArtistCoherencyModel.from_pretrained(
    "tjl223/artist-coherency-ensemble"
)

lyric_generator_model = LyricGeneratorModel(
    "tjl223/testllama2-qlora-lyric-generator-with-description"
)

if submit_button:
    prompt = f"[Song Title] {song_title}\n[Song Artist] {artist_name_input}\n[Song Description] {song_description}"
    lyrics = lyric_generator_model.generate_lyrics(prompt, 1000)
    st.write(lyrics)