import gradio as gr
import os
from huggingface_hub.inference_api import InferenceApi

API_Token = os.environ['HFKEY']
inference = InferenceApi(repo_id="gpt2-large", token=API_Token)

def text_gen(input_string, max_length):
  contrastive_params = {
    "max_length": max_length,
    "penalty_alpha": 0.6,
    "top_k": 4
  }
  contrastive = inference(input_string, contrastive_params)
  cSearchText = contrastive[0]['generated_text']
  return cSearchText
  
def to_gradio():

  demo = gr.Interface(
    fn=text_gen,
    inputs=["text", gr.Slider(0, 250)],
    outputs=["text"],
    title="AI Camp NLP Gradio Mini Project",
    description="Generate text with GPT-2 and contrastive search!",
    css=".gradio-container {background-color:DarkSlateGray}"
  )
  demo.launch()

to_gradio()