Update Similarity.py

Similarity.py

@@ -1,15 +1,8 @@
-import nltk
-nltk.data.path.append("./nltk_data")
-from sentence_transformers import SentenceTransformer, util
+from sentence_transformers import SentenceTransformer
 
 class Similarity:
     def __init__(self):
         self.model = None
-        # Download punkt tokenizer once, suppress if already present
-        try:
-            nltk.data.find('tokenizers/punkt')
-        except LookupError:
-            nltk.download('punkt', download_dir='./nltk_data')
 
     def load_model(self):
         if self.model is None:
@@ -17,43 +10,6 @@ class Similarity:
             self.model = SentenceTransformer("lighteternal/stsb-xlm-r-greek-transfer")
             print("Model loaded.")
 
-    def chunk_text(self, text, chunk_size=1400, overlap_size=200):
-        sentences = nltk.sent_tokenize(text)
-        chunks = []
-        current_chunk = ""
-        for sentence in sentences:
-            if len(current_chunk) + len(sentence) <= chunk_size:
-                current_chunk += " " + sentence if current_chunk else sentence
-            else:
-                chunks.append(current_chunk)
-                # Start the next chunk with overlap
-                current_chunk = sentence[:overlap_size] + sentence[overlap_size:]
-        if current_chunk:
-            chunks.append(current_chunk)
-        return chunks
-
-    def get_sim_text(self, text, claim_embedding, min_threshold=0.4, chunk_size=1500):
+    def embed_text(self, text):
         self.load_model()
-
-        if not text:
-            return []
-
-        filtered_results = []
-        chunks = self.chunk_text(text, chunk_size)
-        if not chunks:
-            return []
-
-        chunk_embeddings = self.model.encode(
-            chunks, convert_to_tensor=True, show_progress_bar=False
-        )
-        chunk_similarities = util.cos_sim(claim_embedding, chunk_embeddings)
-
-        for chunk, similarity in zip(chunks, chunk_similarities[0]):
-            if similarity >= min_threshold:
-                print(chunk)
-                print()
-                print(similarity)
-                print("--------------------------------------------------")
-                filtered_results.append(chunk)
-
-        return filtered_results
+        return self.model.encode(text, convert_to_tensor=True)