Update app.py

app.py

@@ -8,7 +8,7 @@ import io
 
 # Function to perform mean pooling on the model outputs
 def mean_pooling(model_output, attention_mask):
-    token_embeddings = model_output['last_hidden_state']  # Get the embeddings from the model output
+    token_embeddings = model_output['last_hidden_state']
     input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
     sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
     sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
@@ -26,7 +26,7 @@ model_text = AutoModel.from_pretrained('jim33282007/5240_grp27_proj')
 model_gpt2 = AutoModelForCausalLM.from_pretrained('gpt2')
 tokenizer_gpt2 = AutoTokenizer.from_pretrained('gpt2')
 
-st.title('Image Captioning, Text Embedding, and Text Input Application')
+st.title('Image Captioning, Text Embedding, Text Generation, and Input Application')
 
 # Function to load images from URL
 def load_image_from_url(url):
@@ -46,7 +46,6 @@ typed_text = ""
 if input_type == "Upload Image":
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
     if uploaded_file is not None:
-        # Convert bytes to a PIL image
         image = Image.open(io.BytesIO(uploaded_file.getvalue()))
         st.image(image, caption='Uploaded Image', use_column_width=True)
 elif input_type == "Image URL":
@@ -58,26 +57,28 @@ elif input_type == "Image URL":
 elif input_type == "Type Sentence":
     typed_text = st.text_area("Type your sentence here:")
 
-# Generate caption button
-if st.button('Generate Caption and Process'):
+# Generate caption and process text button
+if st.button('Generate Caption and Process Text'):
     if image or typed_text:
         with st.spinner("Processing..."):
             generated_text_p1 = ""
             if input_type == "Upload Image" and uploaded_file is not None:
-                # Use the PIL image directly with the pipeline
                 result = image_to_text(image)
                 generated_text_p1 = result[0]['generated_text']
             elif input_type == "Image URL" and image_url:
                 result = image_to_text(image_url)
                 generated_text_p1 = result[0]['generated_text']
             elif input_type == "Type Sentence" and typed_text:
-                generated_text_p1 = typed_text  # Use the typed text directly
+                generated_text_p1 = typed_text
             
             if generated_text_p1:
                 st.success(f'Processed Text: {generated_text_p1}')
 
-                # Further processing of the text can be done here
-                # This could include generating embeddings, further text, etc., similar to what is done for the image captions
+                # Generate additional text using GPT-2 based on the processed text
+                input_ids = tokenizer_gpt2.encode(generated_text_p1, return_tensors='pt')
+                generated_outputs = model_gpt2.generate(input_ids, max_length=100, num_return_sequences=1)
+                generated_text = tokenizer_gpt2.decode(generated_outputs[0], skip_special_tokens=True)
+                
+                st.text_area("Generated Text:", generated_text, height=200)
     else:
         st.error("Please upload an image, enter an image URL, or type a sentence first.")
-