tareeb23 commited on
Commit
0b2c02e
·
verified ·
1 Parent(s): 3e876e3

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +37 -0
app.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
3
+ import streamlit as st
4
+
5
+ def generate_blog(title, model_name='gpt2', max_length=500):
6
+ # Check if a GPU is available
7
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
8
+ st.write(f"Using device: {device}")
9
+
10
+ # Load the tokenizer and model
11
+ tokenizer = GPT2Tokenizer.from_pretrained(model_name)
12
+ model = GPT2LMHeadModel.from_pretrained(model_name).to(device)
13
+
14
+ # Prepare the input
15
+ input_ids = tokenizer.encode(title, return_tensors='pt').to(device)
16
+
17
+ # Generate text
18
+ output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
19
+
20
+ # Decode the generated text
21
+ blog_post = tokenizer.decode(output[0], skip_special_tokens=True)
22
+
23
+ return blog_post
24
+
25
+ st.title("AI Blog Writer")
26
+ st.write("Enter a blog title, and the AI will generate a blog post for you!")
27
+
28
+ title = st.text_input("Enter the blog title:")
29
+
30
+ if st.button("Generate Blog"):
31
+ if title:
32
+ with st.spinner("Generating blog post..."):
33
+ blog_post = generate_blog(title)
34
+ st.subheader("Generated Blog Post")
35
+ st.write(blog_post)
36
+ else:
37
+ st.warning("Please enter a blog title.")