import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# Load the GPT-2 model and tokenizer
model_name = "gpt2"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model and tokenizer, move model to the correct device
model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Define the sentence completion function
def complete_sentence(sentence):
    if not sentence.strip():
        return "Please enter a valid input."
    try:
        # Encode the input sentence
        input_ids = tokenizer.encode(sentence, return_tensors="pt").to(device)
        
        # Generate completion
        output = model.generate(
            input_ids,
            max_length=50,
            num_return_sequences=1,
            no_repeat_ngram_size=2,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
        )
        
        # Decode the generated sentence
        completed_sentence = tokenizer.decode(output[0], skip_special_tokens=True)
        return completed_sentence
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=complete_sentence,
    inputs="text",
    outputs="text",
    title="Sentence Completion",
    description="Enter a sentence to complete.",
    examples=["I love to", "The future of AI is", "Once upon a time"],
)

# Launch the Gradio interface
if __name__ == "__main__":
    iface.launch()