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 = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"} #TOKEN HUGGING FACE
API_URL = "https://api-inference.huggingface.co/models/BenBranyon/tinyllama-sumbot-peft"

def query(payload):
    json_body = {
        "inputs": f"<|system|> As an AI, you merge the realms of rap, technology, and social activism, embodying the spirit of SUM, a multi-disciplinary, award-winning artist with a foundation in writing 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 Palestine, Congo, Sudan, and 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 content that promotes awareness and human evolution, utilizing humor and a distinctive voice to connect deeply and honor humanity. Each of your responses should be rap lyrics that rhyme like poetry.</s><|user|>{payload}</s><|assistant|>",
        "parameters": {"max_new_tokens":250, "top_p":0.9, "temperature":0.7}
    }
    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 in the style of the artist SUM about {prompt}"
    output = query(prompt)
    st.session_state.messages.append({"role": "assistant", "content": output})
    with st.chat_message("assistant"):
        st.markdown(output)