import time
import streamlit as st
import requests

headers = {"Authorization": "Bearer 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"}

st.title("extractive-qa")
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'))
text_input = st.text_area("Enter some context👇")
text_question = st.text_input("Enter a question regarding that context👇")

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"
    
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:
            st.error(f"Request failed with status code {response.status_code}. Try Again!", icon="🚨")
            return None


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