import gradio as gr
import logging
# from transformers import GPTJForCausalLM, GPT2Tokenizer

# # Load the GPT-J model and tokenizer
# model_name = "EleutherAI/gpt-j-6B"
# tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# model = GPTJForCausalLM.from_pretrained(model_name)


from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load GPT-2 model and tokenizer
#model_name = "gpt2"  # You can use "gpt2-medium" or "gpt2-large" for more power
model_name = "../custom_model/custom_gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

model.config.eos_token_id = tokenizer.eos_token_id
tokenizer.pad_token = tokenizer.eos_token

# Set the pad_token to eos_token


# Function to generate text based on the user input
def generate_text(prompt):
    # Tokenizing the input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=False, padding=False, max_length=512)
    
    # Generate output
    outputs = model.generate(inputs['input_ids'],max_length = 150, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, 
    		top_p=0.95,
    		eos_token_id=model.config.eos_token_id
    	)
    
    # Decode the output
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return generated_text.strip()


# Gradio interface setup
iface = gr.Interface(fn=generate_text, 
                     inputs=gr.inputs.Textbox(lines=10, placeholder="Enter your prompt here..."), 
                     outputs="text")

# Launch the Gradio interface
iface.launch(share=True)