import gradio as gr
from transformers import pipeline

# Load Model
model_name = "AventIQ-AI/gpt2-next-word-prediction"
predictor = pipeline("text-generation", model=model_name)

def predict_next_word(prompt):
    result = predictor(prompt, max_length=len(prompt.split()) + 1, num_return_sequences=1)
    return result[0]['generated_text']

# Examples
examples = [
    ["Artificial intelligence is"],
    ["The future of technology"],
    ["Machine learning enables"],
    ["Deep learning models are"],
]

# Gradio Interface
def main():
    with gr.Blocks(theme="soft") as demo:
        gr.Markdown("""
        # 🚀 Next-Word Prediction
        Enter a partial sentence, and the model will predict the next word.
        """)
        
        with gr.Row():
            input_text = gr.Textbox(label="Enter a sentence", placeholder="Type here...")
        
        predict_btn = gr.Button("🔮 Predict Next Word")
        output_text = gr.Textbox(label="Predicted Sentence", interactive=False)
        
        predict_btn.click(predict_next_word, inputs=input_text, outputs=output_text)
        
        gr.Examples(examples, inputs=input_text)
        
    demo.launch()

if __name__ == "__main__":
    main()