From 027bb54d1b91a2141ad6dfbd77f2c74b2ba75fad Mon Sep 17 00:00:00 2001
From: Dariush <mohandesdariush@gmail.com>
Date: Fri, 15 Nov 2024 19:19:46 +0100
Subject: [PATCH] improve title matching function

---
 comet/utils/general.py | 32 ++++++++++++++++++++++++++------
 pyproject.toml         |  1 +
 2 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/comet/utils/general.py b/comet/utils/general.py
index 3ff56b6..e9002e3 100644
--- a/comet/utils/general.py
+++ b/comet/utils/general.py
@@ -10,6 +10,7 @@ import orjson
 from RTN import parse, title_match
 from curl_cffi import requests
 from fastapi import Request
+from fuzzywuzzy import fuzz
 
 from comet.utils.logger import logger
 from comet.utils.models import settings, ConfigModel
@@ -465,13 +466,32 @@ async def get_mediafusion(log_name: str, type: str, full_id: str):
 
     return results
 
+def match_titles(imdb_title, torrent_title, threshold=80, token_overlap_threshold=0.5):
+    """
+    Match movie/TV show titles using a combination of fuzzy string matching and token overlap.
+
+    Parameters:
+    imdb_title (str): The title from the IMDB data source.
+    torrent_title (str): The title from the torrent data source.
+    threshold (int): The minimum fuzzy match ratio to consider the titles a match.
+    token_overlap_threshold (float): The minimum proportion of overlapping tokens to consider the titles a match.
+
+    Returns:
+    bool: True if the titles match, False otherwise.
+    """
+    # Calculate the fuzzy match ratio
+    match_ratio = fuzz.token_set_ratio(imdb_title, torrent_title)
+
+    # Calculate the proportion of overlapping tokens
+    imdb_tokens = set(imdb_title.lower().split())
+    torrent_tokens = set(torrent_title.lower().split())
+    common_tokens = imdb_tokens.intersection(torrent_tokens)
+    token_overlap_ratio = len(common_tokens) / max(len(imdb_tokens), len(torrent_tokens))
+
+    # Check if both the fuzzy match ratio and token overlap ratio meet the thresholds
+    return match_ratio >= threshold and token_overlap_ratio >= token_overlap_threshold
 
 async def filter(torrents: list, name: str, year: int):
-    def title_sub_match(imdb_title: str, torrent_title: str):
-        imdb_title = imdb_title.lower()
-        torrent_title = torrent_title.lower()
-        return imdb_title in torrent_title or torrent_title in imdb_title
-
     results = []
     for torrent in torrents:
         index = torrent[0]
@@ -482,7 +502,7 @@ async def filter(torrents: list, name: str, year: int):
 
         parsed = parse(title)
 
-        if parsed.parsed_title and not (title_match(name, parsed.parsed_title) or title_sub_match(name, parsed.parsed_title)):
+        if parsed.parsed_title and not (match_titles(name, parsed.parsed_title)):
             results.append((index, False))
             continue
 
diff --git a/pyproject.toml b/pyproject.toml
index 222dbe4..543895a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,6 +25,7 @@ aiosqlite = "*"
 jinja2 = "*"
 rank-torrent-name = "*"
 parsett = "*"
+fuzzywuzzy = "*"
 
 
 [tool.poetry.group.dev.dependencies]