Update main.py

main.py

@@ -3,7 +3,7 @@ from starlette.applications import Starlette
 from starlette.routing import Route
 from starlette.middleware import Middleware
 from starlette.middleware.cors import CORSMiddleware
-
+from gensim.models import KeyedVectors
 """Prompt templates for LLM"""
 from env import LLM_API_KEY
 import prompt
@@ -544,6 +544,45 @@ async def get_flashcards(request):
 
     return JSONResponse({"tldr": __tldr, "defs": __definitions, "imgs": await fetch_img_for_words(__keywords)})
 
+"""
+Similarity validation
+"""
+w2v_vi = KeyedVectors.load_word2vec_format('wiki.vi.model.bin', binary=True)
+# w2v_en = KeyedVectors.load_word2vec_format('GoogleNews-vectors-negative300.bin',binary=True)
+vocab_vi = w2v_vi.key_to_index
+# vocab_en = w2v_en.vocab
+from underthesea import word_tokenize
+from nltk.tokenize import word_tokenize as word_tokenize_en
+from numpy import zeros,zeros_like
+from scipy.spatial.distance import cosine
+
+async def validate_similarity(request):
+    req = await request.json()
+    sent1, sent2 = req['sentences']
+    l = req['lang']
+
+    if (l == lang.VI_VN):
+        tokens1 = word_tokenize(sent1.lower())
+        tokens2 = word_tokenize(sent2.lower())
+    else:
+        tokens1 = word_tokenize_en(sent1.lower())
+        tokens2 = word_tokenize_en(sent2.lower())
+
+    vect1 = zeros_like(w2v_vi.get_vector('an'))
+    vect2 = zeros_like(w2v_vi.get_vector('an'))
+
+    for t in tokens1:
+        if t in vocab_vi:
+            vect1 += w2v_vi.get_vector(t)
+    
+    
+    for t in tokens2:
+        if t in vocab_vi:
+            vect2 += w2v_vi.get_vector(t)
+
+    # Calculate similarity using cosine similarity: This metric measures the cosine of the angle between two embedding vectors. A higher cosine similarity indicates more similar sentences.
+    sim = 1 - cosine(vect1, vect2) >= 0.8
+    return JSONResponse({"isSimilar": str(sim)})
 
 app = Starlette(debug=True,routes=[
     Route('/getFlashcards/{id}/{lang}', get_flashcards, methods=['GET']),
@@ -556,7 +595,7 @@ app = Starlette(debug=True,routes=[
     Route('/generateQuiz/{id}/{lang}', generate_questions, methods=['GET']),
     Route('/convert2md', convert2md, methods=['POST']),
     Route('/mltest', __mltest, methods=['GET'])
-
+    Route('/validateSimilarity', validate_similarity, methods=['POST'])
 ],
 middleware=middleware)