Update app.py

app.py

@@ -36,10 +36,10 @@ def load_embedding(model):
     
     classes = [e["text"] for e in dataset]
     embeddings_list = {cls: model.encode(cls, convert_to_tensor=True) for cls in classes}
-    return embeddings_list
+    return embeddings_list, classes
 
 def generate_image(caption):
-    embeddings = load_embedding(model)
+    embeddings, classes = load_embedding(model)
     noise_dim = 16
     results = [(util.pytorch_cos_sim(model.encode(caption, convert_to_tensor=True), embeddings[cls]).item(), cls) for cls in classes]
     sorted_results = sorted(results, key=lambda x: x[0], reverse=True)[:5]