import streamlit as st
import requests

#API_URL = "https://g8u06j6fqi4vyi5i.eu-west-1.aws.endpoints.huggingface.cloud"
API_URL = "https://gavdzeg3p9lv0fx0.us-east-1.aws.endpoints.huggingface.cloud"
headers = {
	"Accept" : "application/json",
	"Content-Type": "application/json" 
}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()

def get_completion(prompt, temperature=0.6):
    output = query({
	"inputs": f"{prompt}",
	"parameters": {
        "temperature":temperature,
    }
    })
    #If expected format, output text as usual
    if output and isinstance(output, list) and "generated_text" in output[0]:
        return output[0]["generated_text"]
    # Debugging
    if "error" in output:
        return output["error"]
    return str(output.keys())

def main():
    st.title('UiA ai koordinator\n Llama 3 8b text completion')

    user_input = st.text_area("Enter your prompt:", height=300)
    temperature = st.slider('Select the temperature:', 0.0001, 2.0, 0.6)
    
    with st.spinner('Generating completion...'):
        if st.button('Generate'):
            # Get the model's completion
            completion = get_completion(user_input, temperature)
            st.text_area("Model Completion:", value=completion, height=300, key="2")

if __name__ == '__main__':
    main()