Update app.py

app.py

@@ -1,13 +1,18 @@
 import numpy as np
 from PIL import Image
 import streamlit as st
-from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
+from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel, GPT2Tokenizer, GPT2LMHeadModel
 
 # Directory path to the saved model on Google Drive
 model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 
+# Load the pre-trained model and tokenizer
+model_name = "gpt2"
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
 def generate_captions(image):
     image = Image.open(image).convert("RGB")
     generated_caption = tokenizer.decode(model.generate(feature_extractor(image, return_tensors="pt").pixel_values.to("cpu"))[0])
@@ -16,6 +21,18 @@ def generate_captions(image):
     generated_caption = sentence.replace(text_to_remove, "")
     return generated_caption
 
+# Define the Streamlit app
+def generate_paragraph(prompt):
+    # Tokenize the prompt
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+
+    # Generate the paragraph
+    output = model.generate(input_ids, max_length=200, num_return_sequences=1, early_stopping=True)
+
+    # Decode the generated output into text
+    paragraph = tokenizer.decode(output[0], skip_special_tokens=True)
+    return paragraph
+    
 # create the Streamlit app
 def app():
     st.title('Image from your Side, Trending Hashtags from our Side')
@@ -36,6 +53,9 @@ def app():
         st.image(image, caption='The Uploaded File')
         st.write("First is first captions for your Photo : ", string)
 
+        generated_paragraph = generate_paragraph(string)
+
+        st.write(generated_paragraph)
 # run the app
 if __name__ == '__main__':
     app()