import os
import requests
import gradio as gr
from dotenv import load_dotenv

load_dotenv()

API_URL_FALCON = "https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct"
API_URL_GUANACO = "https://api-inference.huggingface.co/models/timdettmers/guanaco-33b-merged"
API_URL_PYTHIA = "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"

headers = {"Authorization": f"Bearer {os.getenv('HF_API_KEY')}"}

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

def respond(message):
    response_falcon = query(API_URL_FALCON, {"inputs": message})
    response_guanaco = query(API_URL_GUANACO, {"inputs": message})
    response_pythia = query(API_URL_PYTHIA, {"inputs": message})

    generated_text_falcon = response_falcon[0]['generated_text']
    generated_text_guanaco = response_guanaco[0]['generated_text']
    generated_text_pythia = response_pythia[0]['generated_text']

    return generated_text_falcon, generated_text_guanaco, generated_text_pythia

    

iface = gr.Interface(
    respond,
    inputs=gr.inputs.Textbox(label="Prompt for all the different models"),
    outputs=[
        gr.outputs.Textbox(label="Falcon Response"),
        gr.outputs.Textbox(label="Guanaco Response"),
        gr.outputs.Textbox(label="Pythia Response")
    ],
    title = "AI Response Aggregator with different LLM models in HugginFace. 🤗",
    description="The purpose is to show the interaction of different models so you can make rapid comparisons 🖥️💡",
    article="<p>This interface allows users to compare real-time outputs from multiple AI models, namely Falcon, Guanaco, and Pythia. By inputting a prompt, users can observe the different ways each model responds, providing a comprehensive view of their capabilities and styles.</p>"
)

iface.launch()