import json
import copy
import gradio as gr

import tensorflow as tf
import tensorflow_text
from tensorflow.python.saved_model import tag_constants
from huggingface_hub import Repository

from pingpong import PingPong
from pingpong.gradio import GradioAlpacaChatPPManager
from pingpong.context import CtxLastWindowStrategy

local_path = "hf_model"

model_version = "v1687590401"
model_repo_id = "chansung/kerasnlp-gpt2-alpaca-pipeline"
model_repo_url = f"https://huggingface.co/{model_repo_id}"

def _clone_and_checkout(repo_url: str, local_path: str, version: str) -> Repository:
    repository = Repository(
        local_dir=local_path, clone_from=repo_url
    )
    repository.git_checkout(revision=version)
    return repository

_ = _clone_and_checkout(model_repo_url, local_path, model_version)
model = tf.saved_model.load(local_path, tags=[tag_constants.SERVING])
# gpt_lm_predict_fn = model.signatures["serving_default"]

with gr.Blocks() as demo:
    gr.Markdown("Hello")

demo.launch()