Update app.py

app.py

@@ -7,9 +7,16 @@ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoToken
 import nltk
 import tempfile
 import zipfile
+from nltk.corpus import wordnet
+import spacy
+import io
+from spacy.cli import download
 
+# Download necessary NLP models
 nltk.download('wordnet')
 nltk.download('omw-1.4')
+download("en_core_web_sm")
+nlp = spacy.load("en_core_web_sm")
 
 # Load the pre-trained models for image captioning and summarization
 model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09"
@@ -35,15 +42,22 @@ def get_synonyms(word):
             synonyms.add(lemma.name())
     return synonyms
 
+def preprocess_query(query):
+    doc = nlp(query)
+    tokens = set()
+    for token in doc:
+        tokens.add(token.text)
+        tokens.add(token.lemma_)
+        tokens.update(get_synonyms(token.text))
+    return tokens
+
 def search_captions(query, captions):
-    query_words = query.split()
-    query_synonyms = set(query_words)
-    for word in query_words:
-        query_synonyms.update(get_synonyms(word))
+    query_tokens = preprocess_query(query)
     
     results = []
     for path, caption in captions.items():
-        if any(word in caption.split() for word in query_synonyms):
+        caption_tokens = preprocess_query(caption)
+        if query_tokens & caption_tokens:
             results.append((path, caption))
     
     return results