Update app.py

app.py

@@ -1,32 +1,18 @@
+from sentence_transformers import SentenceTransformer, util
 from PIL import Image
 import gradio as gr
 import requests
-from transformers import AutoTokenizer, AutoModel
-
-
 
 def get_image_embedding(image):
-    return {"embedding": "img_emb.tolist()"}
+    return {"embedding": "img_emb.tolist()")}
 
 def get_text_embedding(text):
-        # Load the tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("Alibaba-NLP/gte-Qwen2-1.5B-instruct")
-    
-    # Load the model
-    model = AutoModel.from_pretrained("Alibaba-NLP/gte-Qwen2-1.5B-instruct")
-
-    # Tokenize the input text
-    text = "Your input text goes here"
-    inputs = tokenizer(text, return_tensors='pt')
-    
-    # Get embeddings from the model
-    with torch.no_grad():
-    outputs = model(**inputs)
-    embeddings = outputs.last_hidden_state
-    
-    # Process embeddings (e.g., take the mean of all token embeddings)
-    sentence_embedding = embeddings.mean(dim=1)
-    return {"embedding": sentence_embedding}
+    multilingual_text_model = SentenceTransformer('Alibaba-NLP/gte-Qwen2-1.5B-instruct')
+    text_emb = multilingual_text_model.encode(text)
+    print(text_emb)
+    print(type(text_emb))
+    print(text_emb.ndim)
+    return {"embedding": text_emb.tolist()}
 
 image_embedding = gr.Interface(fn=get_image_embedding, inputs=gr.Image(type="pil"), outputs=gr.JSON(), title="Image Embedding")
 text_embedding = gr.Interface(fn=get_text_embedding, inputs=gr.Textbox(), outputs=gr.JSON(), title="Text Embedding")