import gradio as gr

import os

import requests
headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
def query(payload):
  response = requests.post(API_URL, headers=headers, json=payload)
  return response.json()
        
def yes_man(message, history):
   # if message.endswith("?"):
    #    return "Yes"
    #else:
    #    return "Ask me anything!"
    data = [{'error':'just started'}]
    while 'generated_text' not in list(data[0].keys()):
        data = query({"inputs": message})
    print(data)
    return data[0]['generated_text']
gr.ChatInterface(
    yes_man,
    chatbot=gr.Chatbot(height=300),
    textbox=gr.Textbox(placeholder="Ask me a yes or no question", container=False, scale=7),
    title="Yes Man",
    description="Ask Yes Man any question",
    theme="soft",
    examples=["Hello", "Am I cool?", "Are tomatoes vegetables?"],
    cache_examples=True,
    retry_btn=None,
    undo_btn="Delete Previous",
    clear_btn="Clear",
).launch()