import time
import streamlit as st
import requests
headers = {"Authorization": "Bearer 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"}


st.title("question-answering-demo")
option = st.selectbox('Select a model', ('distilbert-base-cased-distilled-squad', 'roberta-base-squad2-distilled', 'xlm-roberta-large-squad2', 'bert-large-cased-whole-word-masking-finetuned-squad'))
if option == "distilbert-base-cased-distilled-squad":
    API_URL = 'https://api-inference.huggingface.co/models/distilbert-base-cased-distilled-squad'
elif option == 'roberta-base-squad2-distilled':
    API_URL = 'https://api-inference.huggingface.co/models/deepset/roberta-base-squad2-distilled'
elif option == 'xlm-roberta-large-squad2':
    API_URL = "https://api-inference.huggingface.co/models/deepset/xlm-roberta-large-squad2"
if option == 'bert-large-cased-whole-word-masking-finetuned-squad':
    API_URL = 'https://api-inference.huggingface.co/models/bert-large-cased-whole-word-masking-finetuned-squad'
st.write(API_URL)
text_input = st.text_area("Enter some context👇")
text_question = st.text_input("Enter a question regarding that context👇")


def query(payload):
    retries = 0
    while True:
        response = requests.post(API_URL, headers=headers, json=payload)
        if response.status_code == 200:
            return response.json()
        elif response.status_code == 429:
            retries += 1
            wait_time = 2 ** retries
            print(f"Too many requests. Retrying in {wait_time} seconds...")
            time.sleep(wait_time)
        else:
            print(f"Request failed with status code {response.status_code}.")
            return None


if st.button("Send"):
    output = query({"inputs": {"question": text_question, "context": text_input}})
    st.write(output)