|
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) |
|
|
|
|
|
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." |
|
) |
|
|
|
|
|
iface.launch() |
|
|