Update chatbot.py

chatbot.py

@@ -1,7 +1,6 @@
 import os
 import pickle
 import torch
-import pickle
 import matplotlib.pyplot as plt
 from langchain_community.document_loaders import TextLoader
 from datasets import load_dataset
@@ -9,6 +8,8 @@ from sentence_transformers import SentenceTransformer, util
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from transformers import BertModel, BertTokenizer
 from langchain_core.prompts import PromptTemplate
+from transformers import BlipProcessor, BlipForConditionalGeneration
+from PIL import Image
 
 os.environ['HUGGINGFACEHUB_API_TOKEN'] = "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"
 
@@ -45,6 +46,8 @@ class Chatbot:
         self.gpt2_model_name = "gpt2"
         self.gpt2_model = GPT2LMHeadModel.from_pretrained(self.gpt2_model_name)
         self.gpt2_tokenizer = GPT2Tokenizer.from_pretrained(self.gpt2_model_name)
+        self.blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+        self.blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
 
     def load_embeddings(self):
         if os.path.exists("embeddings_cache.pkl"):
@@ -102,11 +105,12 @@ class Chatbot:
             plt.axis('off')
             plt.show()
 
-    @staticmethod
-    def cos_sim(a, b):
-        a_norm = torch.nn.functional.normalize(a, p=2, dim=1)
-        b_norm = torch.nn.functional.normalize(b, p=2, dim=1)
-        return torch.mm(a_norm.T, b_norm)  # Reshape a_norm to (768, 1)
+    def generate_image_caption(self, image_path):
+        raw_image = Image.open(image_path).convert('RGB')
+        inputs = self.blip_processor(raw_image, return_tensors="pt")
+        out = self.blip_model.generate(**inputs)
+        caption = self.blip_processor.decode(out[0], skip_special_tokens=True)
+        return caption
 
     def generate_response(self, query):
         # Process the user query and generate a response
@@ -119,4 +123,4 @@ class Chatbot:
         self.display_text_and_images(results_text)
 
         # Return both chatbot response and recommended products
-        return chatbot_response,results_text
+        return chatbot_response, results_text