Update inference.py

inference.py

@@ -1,4 +1,3 @@
-
 from PIL import Image
 import torch
 from transformers import CLIPProcessor, CLIPModel, T5Tokenizer, T5ForConditionalGeneration
@@ -14,7 +13,12 @@ clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 text_encoder = SentenceTransformer('all-MiniLM-L6-v2')
 tokenizer = T5Tokenizer.from_pretrained("t5-small")
 generator = T5ForConditionalGeneration.from_pretrained("t5-small")
-nlp = spacy.load("en_core_web_sm")
+try:
+    nlp = spacy.load("en_core_web_sm")
+except:
+    import spacy.cli
+    spacy.cli.download("en_core_web_sm")
+    nlp = spacy.load("en_core_web_sm")
 
 # Load FAISS index and captions
 faiss_index = faiss.read_index("./faiss_index.idx")
@@ -28,7 +32,6 @@ def extract_image_features(image):
     Output: Normalized image embedding (numpy array).
     """
     try:
-        # Handle both PIL Image and file path
         if isinstance(image, str):
             image = Image.open(image).convert("RGB")
         else:
@@ -98,13 +101,4 @@ def generate_rag_caption(image):
     retrieved = retrieve_similar_captions(embedding, k=5)
     if not retrieved:
         return "No similar captions found."
-    return generate_caption_from_retrieved(retrieved)
-
-def predict(image):
-    """
-    API-compatible function for inference.
-    Input: PIL Image or image file path.
-    Output: Dictionary with caption.
-    """
-    caption = generate_rag_caption(image)
-    return {"caption": caption}
+    return generate_caption_from_retrieved(retrieved)