Update app.py

app.py

@@ -1,20 +1,32 @@
-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):
-    image_model = SentenceTransformer('clip-ViT-B-32')
-    img_emb = image_model.encode(image)
-    return {"embedding": img_emb.tolist()}
+    return {"embedding": "img_emb.tolist()"}
 
 def get_text_embedding(text):
-    multilingual_text_model = SentenceTransformer('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()}
+        # 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}
 
 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")