src/call_interface.py

"""Call the inference API of HF.
"""

import huggingface_hub

import gradio as gr


model_name = "yp-edu/gpt2-stockfish-debug"

headers = {"X-Wait-For-Model": "true"}
client = huggingface_hub.InferenceClient(
    model=model_name, headers=headers
)

inputs = gr.Textbox(label="Prompt")
outputs = gr.Textbox(label="Completion")
examples = [
    "FEN: rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1\nMOVE:",
    "FEN: r2q1rk1/1p3ppp/4bb2/p2p4/5B2/1P1P4/1PPQ1PPP/R3R1K1 w - - 1 17\nMOVE:",
    "FEN: 4r1k1/1p1b1ppp/8/8/3P4/2P5/1q3PPP/6K1 b - - 0 28\nMOVE:",
]
fn = client.text_generation


interface = gr.Interface(
    fn=fn,
    inputs=inputs,
    outputs=outputs,
    title=f"Call to {model_name}",
    examples=examples,
)