import streamlit as st
import tensorflow as tf
from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
import pandas as pd
import numpy as np

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = TFGPT2LMHeadModel.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id)

def func(sentence, max_length, temperature):
    input_ids = tokenizer.encode(sentence, return_tensors='tf')
    output_list = model.generate(
        input_ids, 
        do_sample=True, 
        max_length=max_length,
        temperature=temperature,
        top_p=0.92, 
        top_k=0,
        num_return_sequences=5
    )
    output_strs = [tokenizer.decode(output, skip_special_tokens=True) for output in output_list]
    return output_strs

sentence = st.text_input(label="Sentence to complete")
max_length = st.slider(label="Max Length", min_value=5, max_value=25, value=10, step=1)
temperature = st.slider(label="Temperature", min_value=0.1, max_value=10.0, value=0.1)

if st.button('Click to generate possible completions'):
    outputs_strs = func(sentence, max_length, temperature)
    i = 1
    for output in outputs_strs:
        st.write(f"{i}: {output}")
        i += 1