Chakshu123's picture
Duplicate from Chakshu123/sketch-colorization-with-hint
443d045
raw
history blame
6.09 kB
"""Utility function for gradio/external.py"""
import base64
import json
import math
import operator
import re
import warnings
from typing import Any, Dict, List, Tuple
import requests
import websockets
import yaml
from packaging import version
from websockets.legacy.protocol import WebSocketCommonProtocol
from gradio import components, exceptions
##################
# Helper functions for processing tabular data
##################
def get_tabular_examples(model_name: str) -> Dict[str, List[float]]:
readme = requests.get(f"https://huggingface.co/{model_name}/resolve/main/README.md")
if readme.status_code != 200:
warnings.warn(f"Cannot load examples from README for {model_name}", UserWarning)
example_data = {}
else:
yaml_regex = re.search(
"(?:^|[\r\n])---[\n\r]+([\\S\\s]*?)[\n\r]+---([\n\r]|$)", readme.text
)
if yaml_regex is None:
example_data = {}
else:
example_yaml = next(
yaml.safe_load_all(readme.text[: yaml_regex.span()[-1]])
)
example_data = example_yaml.get("widget", {}).get("structuredData", {})
if not example_data:
raise ValueError(
f"No example data found in README.md of {model_name} - Cannot build gradio demo. "
"See the README.md here: https://huggingface.co/scikit-learn/tabular-playground/blob/main/README.md "
"for a reference on how to provide example data to your model."
)
# replace nan with string NaN for inference API
for data in example_data.values():
for i, val in enumerate(data):
if isinstance(val, float) and math.isnan(val):
data[i] = "NaN"
return example_data
def cols_to_rows(
example_data: Dict[str, List[float]]
) -> Tuple[List[str], List[List[float]]]:
headers = list(example_data.keys())
n_rows = max(len(example_data[header] or []) for header in headers)
data = []
for row_index in range(n_rows):
row_data = []
for header in headers:
col = example_data[header] or []
if row_index >= len(col):
row_data.append("NaN")
else:
row_data.append(col[row_index])
data.append(row_data)
return headers, data
def rows_to_cols(incoming_data: Dict) -> Dict[str, Dict[str, Dict[str, List[str]]]]:
data_column_wise = {}
for i, header in enumerate(incoming_data["headers"]):
data_column_wise[header] = [str(row[i]) for row in incoming_data["data"]]
return {"inputs": {"data": data_column_wise}}
##################
# Helper functions for processing other kinds of data
##################
def postprocess_label(scores: Dict) -> Dict:
sorted_pred = sorted(scores.items(), key=operator.itemgetter(1), reverse=True)
return {
"label": sorted_pred[0][0],
"confidences": [
{"label": pred[0], "confidence": pred[1]} for pred in sorted_pred
],
}
def encode_to_base64(r: requests.Response) -> str:
# Handles the different ways HF API returns the prediction
base64_repr = base64.b64encode(r.content).decode("utf-8")
data_prefix = ";base64,"
# Case 1: base64 representation already includes data prefix
if data_prefix in base64_repr:
return base64_repr
else:
content_type = r.headers.get("content-type")
# Case 2: the data prefix is a key in the response
if content_type == "application/json":
try:
content_type = r.json()[0]["content-type"]
base64_repr = r.json()[0]["blob"]
except KeyError:
raise ValueError(
"Cannot determine content type returned" "by external API."
)
# Case 3: the data prefix is included in the response headers
else:
pass
new_base64 = "data:{};base64,".format(content_type) + base64_repr
return new_base64
##################
# Helper functions for connecting to websockets
##################
async def get_pred_from_ws(
websocket: WebSocketCommonProtocol, data: str, hash_data: str
) -> Dict[str, Any]:
completed = False
resp = {}
while not completed:
msg = await websocket.recv()
resp = json.loads(msg)
if resp["msg"] == "queue_full":
raise exceptions.Error("Queue is full! Please try again.")
if resp["msg"] == "send_hash":
await websocket.send(hash_data)
elif resp["msg"] == "send_data":
await websocket.send(data)
completed = resp["msg"] == "process_completed"
return resp["output"]
def get_ws_fn(ws_url, headers):
async def ws_fn(data, hash_data):
async with websockets.connect( # type: ignore
ws_url, open_timeout=10, extra_headers=headers
) as websocket:
return await get_pred_from_ws(websocket, data, hash_data)
return ws_fn
def use_websocket(config, dependency):
queue_enabled = config.get("enable_queue", False)
queue_uses_websocket = version.parse(
config.get("version", "2.0")
) >= version.Version("3.2")
dependency_uses_queue = dependency.get("queue", False) is not False
return queue_enabled and queue_uses_websocket and dependency_uses_queue
##################
# Helper function for cleaning up an Interface loaded from HF Spaces
##################
def streamline_spaces_interface(config: Dict) -> Dict:
"""Streamlines the interface config dictionary to remove unnecessary keys."""
config["inputs"] = [
components.get_component_instance(component)
for component in config["input_components"]
]
config["outputs"] = [
components.get_component_instance(component)
for component in config["output_components"]
]
parameters = {
"article",
"description",
"flagging_options",
"inputs",
"outputs",
"theme",
"title",
}
config = {k: config[k] for k in parameters}
return config