import gradio as gr
import requests
import json
import os 


API_KEY = os.getenv('API_KEY') 
headers = {
    "Authorization": f"Bearer {API_KEY}",
    "Accept": "application/json",
}

def generate_embedding(input_text, model_type, encoding_format):
    invoke_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/091a03bb-7364-4087-8090-bd71e9277520"
    payload = {
        "input": input_text,
        "model": model_type,
        "encoding_format": encoding_format
    }

    session = requests.Session()
    response = session.post(invoke_url, headers=headers, json=payload)

    while response.status_code == 202:
        request_id = response.headers.get("NVCF-REQID")
        fetch_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/" + request_id
        response = session.get(fetch_url, headers=headers)

    response.raise_for_status()
    response_body = response.json()
    embedding = response_body["data"][0]["embedding"]
    return json.dumps(embedding, indent=2)

# Criação da interface Gradio
iface = gr.Interface(
    fn=generate_embedding,
    inputs=[
        gr.Textbox(label="Input Text"),
        gr.Radio(choices=["query", "passage"], value="query", label="Model Type"),
        gr.Radio(choices=["float", "base64"], value="float", label="Encoding Format")
    ],
    outputs=[gr.Textbox(label="Embedding")],
    title="NVIDIA Retrieval QA Embedding Generator",
    description="Generate embeddings for your text using NVIDIA's Retrieval QA Embedding model."
)

# Executar a aplicação Gradio
iface.launch()