Update Similarity.py

Similarity.py

@@ -1,47 +1,29 @@
-import nltk
-from sentence_transformers import util
+from fastapi import FastAPI
+from pydantic import BaseModel
+from Similarity import Similarity
 
-class Similarity:
-    def __init__(self):
-        self.model = SentenceTransformer("lighteternal/stsb-xlm-r-greek-transfer")
-        # Make sure nltk punkt tokenizer is downloaded
-        nltk.download('punkt')
+app = FastAPI()
+similarity_model = Similarity()
 
-    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
+class TextPairRequest(BaseModel):
+    text: str     # The big text to chunk and search
+    claim: str    # The claim text to embed and compare
 
-    def get_sim_text(self, text, claim_embedding, min_threshold=0.4, chunk_size=1500):
-        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
+@app.post("/get_sim_text")
+def get_sim_text_endpoint(request: TextPairRequest):
+    try:
+        # Embed the claim text
+        claim_embedding = similarity_model.model.encode(
+            request.claim,
+            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)
+        # Call get_sim_text with defaults (min_threshold=0.4, chunk_size=1500)
+        results = similarity_model.get_sim_text(
+            request.text,
+            claim_embedding
+        )
+        return {"result": results}
 
-        return filtered_results
+    except Exception as e:
+        return {"error": str(e)}