Update app.py

app.py

@@ -9,10 +9,12 @@ from PIL import Image
 from transformers import CLIPProcessor, CLIPModel
 from datasets import load_dataset
 
+model_checkpoint = "openai/clip-vit-base-patch32"
+
 def get_clip_embeddings(input_data, input_type='text'):
     # Load the CLIP model and processor
-    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+    model = CLIPModel.from_pretrained(model_checkpoint)
+    processor = CLIPProcessor.from_pretrained(model_checkpoint)
 
     # Prepare the input based on the type
     if input_type == 'text':
@@ -42,17 +44,16 @@ veggies = load_dataset('vojtam/vegetables')
 with open('img_embeddings.pkl', 'rb') as file:
     img_embeddings = pickle.load(file)
 
-text = gr.Textbox(label = "Enter the text", 'Your text goes here')
+cos = nn.CosineSimilarity(dim=1, eps=1e-6)
+
+text = gr.Textbox(label = "Enter the description of the images you want to search for", placeholder='Your text goes here')
 image = gr.Gallery()
 
 def get_similar_images(text, n = 4):
     if text:
         text_embedding = get_clip_embeddings(text, input_type='text')
-        cos = nn.CosineSimilarity(dim=1, eps=1e-6)
         sims = cos(torch.tensor(text_embedding), torch.tensor(img_embeddings))
         top_n = np.argsort(np.array(sims))[::-1][:n]
-        print(top_n)
-        print(img_embeddings)
         imgs = []
         
         for index in top_n: