Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -13,41 +13,11 @@ from fastapi import FastAPI, Body, Request
 from utils.utils import add_arguments, print_arguments
 from sentence_transformers import SentenceTransformer, models
 
+from gensim.models import Word2Vec
+from gensim.utils import simple_preprocess
+import numpy as np
 
 
-# def print_arguments(args):
-#     print("-----------  Configuration Arguments -----------")
-#     for arg, value in vars(args).items():
-#         print("%s: %s" % (arg, value))
-#     print("------------------------------------------------")
-
-
-# def strtobool(val):
-#     val = val.lower()
-#     if val in ('y', 'yes', 't', 'true', 'on', '1'):
-#         return True
-#     elif val in ('n', 'no', 'f', 'false', 'off', '0'):
-#         return False
-#     else:
-#         raise ValueError("invalid truth value %r" % (val,))
-
-# def str_none(val):
-#     if val == 'None':
-#         return None
-#     else:
-#         return val
-
-# def add_arguments(argname, type, default, help, argparser, **kwargs):
-#     type = strtobool if type == bool else type
-#     type = str_none if type == str else type
-#     argparser.add_argument(
-#         "--" + argname,
-#         default=default,
-#         type=type,
-#         help=help + ' Default: %(default)s.',
-#         **kwargs
-#     )
-
 
 os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
@@ -60,11 +30,8 @@ add_arg("host",        type=str,  default="0.0.0.0", help="")
 add_arg("port",        type=int,  default=5000,        help="")
 add_arg("model_path",  type=str,  default="BAAI/bge-small-en-v1.5", help="")
 add_arg("use_gpu",     type=bool, default=False,   help="")
-# add_arg("use_int8",    type=bool, default=True,  help="")
-add_arg("beam_size",   type=int,  default=10,     help="")
 add_arg("num_workers", type=int,  default=2,      help="")
-add_arg("vad_filter",  type=bool, default=True,  help="")
-add_arg("local_files_only", type=bool, default=True, help="")
+
 
 
 args = parser.parse_args()
@@ -72,13 +39,24 @@ print_arguments(args)
 
 
 
+# similarity score func
+def similarity_score(model, textA, textB):
+    em_test = model.encode(
+        [textA, textB],
+        normalize_embeddings=True
+    )
+    return em_test[0] @ em_test[1].T
+
+
+# BGE embedding
+
 if args.use_gpu:
     bge_model = SentenceTransformer(args.model_path, device="cuda", compute_type="float16", cache_folder=".")
 else:
     bge_model = SentenceTransformer(args.model_path, device='cpu', cache_folder=".")
 
 
-
+# tsdae embedding
 if args.use_gpu:
     model_name = 'sam2ai/sbert-tsdae'
     word_embedding_model = models.Transformer(model_name)
@@ -100,22 +78,34 @@ else:
     )
 
 
-app = FastAPI(title="embedding Inference")
+# word2vec embedding
+# Define the calculate_similarity function
+def calculate_similarity(sentence1, sentence2):
+    # Tokenize the sentences
+    tokens1 = simple_preprocess(sentence1)
+    tokens2 = simple_preprocess(sentence2)
 
+    # Load or train a Word2Vec model
+    # Here, we'll create a simple model for demonstration purposes
+    sentences = [tokens1, tokens2]
+    model = Word2Vec(sentences, vector_size=100, window=5, min_count=1, sg=0)
 
-def similarity_score(model, textA, textB):
-    em_test = model.encode(
-        [textA, textB],
-        normalize_embeddings=True
-    )
-    return em_test[0] @ em_test[1].T
+    # Calculate the vector representation for each sentence
+    vector1 = np.mean([model.wv[token] for token in tokens1], axis=0)
+    vector2 = np.mean([model.wv[token] for token in tokens2], axis=0)
+
+    # Calculate cosine similarity
+    similarity = np.dot(vector1, vector2) / (np.linalg.norm(vector1) * np.linalg.norm(vector2))
+    return similarity
 
 
+
+app = FastAPI(title="embedding Inference")
+
 @app.get("/")
 async def index(request: Request):
     return {"detail": "API is Active !!"}
 
-
 @app.post("/bge_embed")
 async def api_bge_embed(
             text1: str = Body("text1", description="", embed=True),
@@ -142,6 +132,19 @@ async def api_tsdae_embed(
     ret = {"similarity score": scores, "status_code": 200}
     return ret
 
+@app.post("/w2v_embed")
+async def api_w2v_embed(
+            text1: str = Body("text1", description="", embed=True),
+            text2: str = Body("text2", description="", embed=True),
+        ):
+
+    scores = calculate_similarity(text1, text2)
+    print(scores)
+    scores = scores.tolist()
+
+    ret = {"similarity score": scores, "status_code": 200}
+    return ret
+