Update app.py

app.py

@@ -1,76 +1,40 @@
 import streamlit as st
-from datasets import load_dataset, DatasetDict
-from transformers import GPT2Tokenizer, GPT2LMHeadModel, DataCollatorForLanguageModeling, Trainer, TrainingArguments
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-@st.cache_resource
-def load_and_fine_tune_model():
-    # Load the dataset
-    dataset = load_dataset("blog_authorship_corpus")
+# Load pre-trained GPT-2 model and tokenizer
+model_name = 'gpt2'
+model = GPT2LMHeadModel.from_pretrained(model_name)
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
-    # Split the dataset into train and validation sets
-    dataset = dataset['train'].train_test_split(test_size=0.1)
-    
-    # Load the tokenizer
-    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-
-    def tokenize_function(examples):
-        return tokenizer(examples["text"], truncation=True)
-
-    tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
-
-    # Data collator
-    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
-
-    # Load the model
-    model = GPT2LMHeadModel.from_pretrained("gpt2")
-
-    # Training arguments
-    training_args = TrainingArguments(
-        output_dir="./results",
-        overwrite_output_dir=True,
-        num_train_epochs=1,
-        per_device_train_batch_size=2,
-        save_steps=10_000,
-        save_total_limit=2,
-    )
+# Function to generate blog post for a given topic
+def generate_blog_post(topic, max_length=300):
+    # Encode the input topic into tokens
+    input_ids = tokenizer.encode(topic, return_tensors='pt')
 
-    # Initialize the Trainer
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        data_collator=data_collator,
-        train_dataset=tokenized_datasets['train'],
-        eval_dataset=tokenized_datasets['test']
-    )
-
-    # Fine-tune the model
-    trainer.train()
-
-    # Save the fine-tuned model
-    model.save_pretrained("./fine-tuned-gpt2")
-    tokenizer.save_pretrained("./fine-tuned-gpt2")
-    return model, tokenizer
-
-def generate_blog_post(prompt, model, tokenizer, max_length=500, temperature=0.7, top_k=50):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    # Generate text
     output = model.generate(
         input_ids,
         max_length=max_length,
-        temperature=temperature,
-        top_k=top_k,
+        num_return_sequences=1,
         no_repeat_ngram_size=2,
-        num_return_sequences=1
+        early_stopping=True
     )
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return generated_text
 
-# Streamlit UI
+    # Decode the output tokens into a string
+    blog_post = tokenizer.decode(output[0], skip_special_tokens=True)
+    
+    return blog_post
+
+# Streamlit app
 st.title("Blog Post Generator")
-prompt = st.text_input("Enter a prompt for the blog post:", "The future of artificial intelligence in daily life")
+
+topic = st.text_input("Enter a topic for the blog post:")
+max_length = st.slider("Maximum length of the blog post:", min_value=50, max_value=1000, value=300)
 
 if st.button("Generate Blog Post"):
-    with st.spinner("Fine-tuning the model. This might take a few minutes..."):
-        model, tokenizer = load_and_fine_tune_model()
-    blog_post = generate_blog_post(prompt, model, tokenizer)
-    st.subheader("Generated Blog Post")
-    st.write(blog_post)
+    if topic:
+        with st.spinner('Generating blog post...'):
+            blog_post = generate_blog_post(topic, max_length)
+        st.write(blog_post)
+    else:
+        st.warning("Please enter a topic.")