import streamlit as st
import os
import requests
import json

#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 art, 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.</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"))
        response = data[0]['generated_text'].split('<|assistant|> ')[1]
        return response
    except:
        return response
    #response = inference(payload)
    #response = requests.post(API_URL, headers=headers, json=payload)
    #return response.json()

messages = st.container()
message = st.chat_message("assistant").write("Hello human")
if prompt := st.chat_input("Subject of the song"):
    messages.chat_message("user").write(prompt)
    prompt = f"Write a rap in the style of the artist Sum about {prompt}"
    output = query(prompt)
    messages.chat_message("assistant").write(output)