From 54ba27ec5a5287f9fb804b51b8f0d4974a78dbb4 Mon Sep 17 00:00:00 2001
From: Brivaldo Junior <condector@gmail.com>
Date: Fri, 11 Apr 2025 17:28:53 -0400
Subject: [PATCH] add support to ollama translation.

---
 deep_translator/__init__.py |  2 ++
 deep_translator/ollama.py   | 59 +++++++++++++++++++++++++++++++++++++
 pyproject.toml              |  1 +
 3 files changed, 62 insertions(+)
 create mode 100644 deep_translator/ollama.py

diff --git a/deep_translator/__init__.py b/deep_translator/__init__.py
index 2b0c9e0..f5c1a65 100644
--- a/deep_translator/__init__.py
+++ b/deep_translator/__init__.py
@@ -16,6 +16,7 @@
 from deep_translator.qcri import QcriTranslator
 from deep_translator.tencent import TencentTranslator
 from deep_translator.yandex import YandexTranslator
+from deep_translator.ollama import OllamaTranslator
 
 __author__ = """Nidhal Baccouri"""
 __email__ = "nidhalbacc@gmail.com"
@@ -35,6 +36,7 @@
     "ChatGptTranslator",
     "TencentTranslator",
     "BaiduTranslator",
+    "OllamaTranslator",
     "single_detection",
     "batch_detection",
 ]
diff --git a/deep_translator/ollama.py b/deep_translator/ollama.py
new file mode 100644
index 0000000..2405df0
--- /dev/null
+++ b/deep_translator/ollama.py
@@ -0,0 +1,59 @@
+__copyright__ = "Copyright (C) 2025 Brivaldo Junior"
+
+from typing import List, Optional
+from deep_translator.base import BaseTranslator
+
+
+class OllamaTranslator(BaseTranslator):
+    """
+    class that wraps functions, which use the Ollama
+    under the hood to translate word(s)
+    """
+
+    def __init__(
+        self,
+        source: str = "auto",
+        target: str = "english",
+        model: Optional[str] = "llama3.2:latest",
+        **kwargs,
+    ):
+        """
+        @param source: source language
+        @param target: target language
+        """
+        self.model = model
+
+        super().__init__(source=source, target=target, **kwargs)
+
+    def translate(self, text: str, **kwargs) -> str:
+        """
+        @param text: text to translate
+        @return: translated text
+        """
+        import ollama
+
+
+        prompt = f"Translate the text below into {self.target}.\n"
+        prompt += f'Text: "{text}"'
+
+        response = ollama.chat(
+            model=self.model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+        )
+
+        return response.message.content
+
+    def translate_file(self, path: str, **kwargs) -> str:
+        return self._translate_file(path, **kwargs)
+
+    def translate_batch(self, batch: List[str], **kwargs) -> List[str]:
+        """
+        @param batch: list of texts to translate
+        @return: list of translations
+        """
+        return self._translate_batch(batch, **kwargs)
diff --git a/pyproject.toml b/pyproject.toml
index fb77272..cfe7914 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,6 +59,7 @@ requests = "^2.23.0"
 docx2txt = {version = "^0.8", optional = true}
 pypdf = {version = "^3.3.0", optional = true}
 openai = {version = "^0.27.6", python = "^3.7.1", optional = true}
+ollama = {version = "^0.4.7", python = "^3.7.1", optional = true}
 
 [tool.poetry.extras]
 docx = ["docx2txt"]