Update app.py

app.py

@@ -1,13 +1,24 @@
 import streamlit as st
-from transformers import pipeline as transformers_pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModel, pipeline as transformers_pipeline, AutoModelForCausalLM
 from diffusers import DiffusionPipeline
 import requests
 from PIL import Image
 import io
+import torch
+import torch.nn.functional as F
+import pandas as pd
 
-# Load the tokenizer and model for similarity checking
-similarity_tokenizer = AutoTokenizer.from_pretrained("AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27")
-similarity_model = AutoModelForSequenceClassification.from_pretrained("AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27")
+# Function for mean pooling of embeddings
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0]  # First element of model_output contains all token embeddings
+    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)
+    return sum_embeddings / sum_mask
+
+# Load model and tokenizer from HuggingFace Hub for sentence embeddings
+tokenizer = AutoTokenizer.from_pretrained('AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27')
+model = AutoModel.from_pretrained('AidenYan/MiniLM_L6_v2_finetuned_ISOM5240_Group27')
 
 def load_image(input_type, uploaded_file=None, image_url=""):
     """
@@ -38,22 +49,37 @@ def select_closest_sentence(generated_text):
     """
     Selects the sentence closest in meaning to the generated_text.
     """
-    # Implement your logic for sentence similarity here.
-    # This is a placeholder function. You need to implement the logic
-    # based on your requirements and the similarity model's capabilities.
+    # Load CSV data
+    df = pd.read_csv('toys_and_games_reviews.csv', encoding='ISO-8859-1')
+    sentences = df.iloc[:, -1].tolist()  # Assuming the last column contains sentences
+
+    # Tokenize and compute embeddings for sentences from CSV
+    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+    with torch.no_grad():
+        model_output = model(**encoded_input)
+    sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
+    sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
     
-    # For the sake of this example, let's return the generated_text itself.
-    return generated_text
+    # Tokenize and compute embedding for the generated_text
+    encoded_new_sentence = tokenizer([generated_text], padding=True, truncation=True, return_tensors='pt')
+    with torch.no_grad():
+        model_output_new_sentence = model(**encoded_new_sentence)
+    new_sentence_embedding = mean_pooling(model_output_new_sentence, encoded_new_sentence['attention_mask'])
+    new_sentence_embedding = F.normalize(new_sentence_embedding, p=2, dim=1)
+
+    # Find the most similar sentence in your corpus
+    most_similar_idx = F.cosine_similarity(new_sentence_embedding, sentence_embeddings).topk(1).indices.item()
+    most_similar_sentence = sentences[most_similar_idx]
+
+    return most_similar_sentence
 
 def generate_text_from_caption(caption):
     """
     Generates text based on the provided caption.
     """
-    model = AutoModelForCausalLM.from_pretrained('pranavpsv/genre-story-generator-v2')
-    tokenizer = AutoTokenizer.from_pretrained('pranavpsv/genre-story-generator-v2')
-    input_ids = tokenizer.encode(caption, return_tensors='pt')
-    generated_outputs = model.generate(input_ids, max_length=100, num_return_sequences=1)
-    return tokenizer.decode(generated_outputs[0], skip_special_tokens=True)
+    text_generator = transformers_pipeline('text-generation', model='pranavpsv/genre-story-generator-v2')
+    generated = text_generator(caption, max_length=100, num_return_sequences=1)
+    return generated[0]['generated_text']
 
 def main():
     st.title('Image to Story to Image Converter')
@@ -61,36 +87,4 @@ def main():
     # User interface for input selection
     input_type = st.radio("Select input type:", ("Upload Image", "Image URL"))
     uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) if input_type == "Upload Image" else None
-    image_url = st.text_input("Enter the image URL here:", "") if input_type == "Image URL" else ""
-
-    # Load image based on input selection
-    image = load_image(input_type, uploaded_file, image_url)
-    if image:
-        st.image(image, caption='Selected Image', use_column_width=True)
-
-    # Process image and generate text
-    if st.button('Generate Caption and Continue'):
-        if image:
-            with st.spinner("Processing..."):
-                # Convert image to text
-                caption = image_to_caption(image, input_type, uploaded_file, image_url)
-                st.success(f'Caption: {caption}')
-
-                # Select the closest sentence using the similarity model
-                closest_sentence = select_closest_sentence(caption)
-                st.write(f"Selected Sentence: {closest_sentence}")
-
-                # Generate additional text based on the selected sentence
-                generated_text = generate_text_from_caption(closest_sentence)
-                st.text_area("Generated Story:", generated_text, height=200)
-
-                # Generate an image from the story
-                diffusion_pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
-                generated_images = diffusion_pipeline(generated_text, num_inference_steps=50)
-                st.image(generated_images.images[0], caption='Generated Image from Story')
-
-        else:
-            st.error("Please upload an image or enter an image URL first.")
-
-if __name__ == "__main__":
-    main()
+    image_url = st.text_input("Enter the image URL