Update app.py

app.py

@@ -1,11 +1,23 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import random
+from datasets import load_dataset
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import torch
 
-# 1️⃣ Load Model & Tokenizer from Hugging Face Hub
+# Load dataset
+dataset = load_dataset("rahul7star/hindi-poetry")["train"]
+
+# Load your model and tokenizer
 model_name = "rahul7star/hindi_poetry_language_model"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+# 1️⃣ Load Model & Tokenizer from Hugging Face Hub
+# model_name = "rahul7star/hindi_poetry_language_model"
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Ensure model uses the correct pad token
 model.config.pad_token_id = tokenizer.pad_token_id