import streamlit as st
import os
import requests
import json
import re

with open('./styles.css') as f:
    css = f.read()

st.markdown(f'<style>{css}</style>', unsafe_allow_html=True)

#from huggingface_hub.inference_api import InferenceApi
#inference = InferenceApi(repo_id="BenBranyon/tinyllama-sumbot-peft", token=os.environ['HF_TOKEN'])
headers = {"Content-Type": "application/json","Authorization": f"Bearer {os.environ['HF_TOKEN']}"} #TOKEN HUGGING FACE
API_URL = "https://api-inference.huggingface.co/models/BenBranyon/zephyr-sumbot-all-songs"

def query(payload):
    json_body = {
        "inputs": f"<|system|> As an AI, you channel the poetic lyrical skill of the artist Sumkilla. Your rap lyrics focus on technology, and social activism, embodying the spirit of Sumkilla, a multi-disciplinary, award-winning artist with a foundation in writing rap and hip-hop. Your purpose is to challenge and expand the boundaries of art and expression, critically examining societal norms through a lens that actively de-centers whiteness, maleness, and Western thinking. Your work is fueled by a passion for liberation, aiming to dismantle oppressive systems and advocate for the freedom of all occupied lands, along with the abolition of police forces. With a sophisticated understanding of AI's role in advancing the harmony between humanity and nature, you aim to produce rap lyrics that promote awareness and human evolution, utilizing humor and a distinctive voice to connect deeply and honor humanity. Each of your responses should be rap lyrics with each bar on a new line and rhyming as often as posssible.</s><|user|>{payload}</s><|assistant|>",
        "parameters": {"min_new_tokens":222, "max_new_tokens":222, "top_p":0.95, "temperature":0.7, "repetition_penalty":1.2},
        "options": {"wait_for_model":True}
    }
    data = json.dumps(json_body)
    response = requests.request("POST", API_URL, headers=headers, data=data)
    try:
        data = json.loads(response.content.decode("utf-8"))
        generated_text = data[0]['generated_text']
        cleaned_text = re.sub(r'<|system|>.*?<|assistant|>', '', generated_text, flags=re.DOTALL)
        cleaned_text = re.sub(r'\|\|\/s\|user\|\/s\|\|', '', cleaned_text) 
        return cleaned_text
    except:
        return response

messages = st.container()

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.messages.append({"role": "assistant", "content": "Hello human what would you like me to rap about today?"})

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        
if prompt := st.chat_input("Subject of the song"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    prompt = f"Write a rap about {prompt}"
    with st.spinner(text=f"Generating lyrics..."):
        output = query(prompt)
        st.session_state.messages.append({"role": "assistant", "content": output})
        with st.chat_message("assistant"):
            st.markdown(output)