Update app.py

app.py

@@ -12,7 +12,13 @@ MODELS = {'enro': 'BlackKakapo/opus-mt-en-ro',
           'roen': 'BlackKakapo/opus-mt-ro-en',
           'gemma': 'Gargaz/gemma-2b-romanian-better',
           'paraphrase': 'tuner007/pegasus_paraphrase'}
-EMBEDDING_MODELS = ["all-MiniLM-L6-v2", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"]
+EMBEDDING_MODELS = {"all-MiniLM-L6-v2":384,
+"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":384,
+"sentence-transformers/distiluse-base-multilingual-cased-v2":512,
+"sentence-transformers/stsb-xlm-r-multilingual":768,
+"sentence-transformers/use-cmlm-multilingual":768,
+"sentence-transformers/paraphrase-multilingual-mpnet-base-v2":768}
+EMBEDDING_MODEL = "sentence-transformers/distiluse-base-multilingual-cased-v2"
 
 @app.get("/")
 def index(request: Request):
@@ -113,10 +119,11 @@ def bergamot(input_text: list[str] = Query(description="Input list of strings"),
     return {"input": input_text, "translated_text": response, "message_text": message_text}
 
 @app.get("/embed", operation_id="get_embeddings", description="Embed text", tags=["embed"], summary="Embed text")
-def embed(text: str, model: str  = EMBEDDING_MODELS[1]):
+def embed(text: str, model: str  = EMBEDDING_MODEL):
     model = SentenceTransformer(model)
     embeddings = model.encode(text)
-    print(embeddings.shape)
+    print(embeddings.shape, len(embeddings))
+    # similarities = model.similarity(embeddings, embeddings)
     return {"input": text, "embeddings": embeddings.tolist(), "shape": embeddings.shape}
 
 # Create an MCP server based on this app