Update app.py

app.py

@@ -82,8 +82,69 @@ def translate(input_text: str, model_name: str = MODELS['mttcbig'], sl: str = 'e
         translation: str = src.Translate.gemma_direct(input_text, model_name)
     return {"input_text": input_text, "translation": translation, "model_name": model_name, "message": message}
 
-# Keep track of installed (src, tgt) pairs
-installed_pairs = set()
+class Bergamot:
+    import bergamot
+    config = bergamot.ServiceConfig(numWorkers=4)
+    service = bergamot.Service(config)
+        
+    def __init__(self, input_text, sl, tl, model_name):
+        # Keep track of installed (src, tgt) pairs
+        self.installed_pairs = set()
+        self.input_text, self.sl, self.tl, self.model_name = input_text, sl, tl, model_name
+        self.repo_id = "TiberiuCristianLeon/Bergamot"
+        self.branches = ['base', 'base-memory', 'tiny']
+        self.subfolder = f"{sl}{tl}"
+        self.localfolder = f"{subfolder}/{model_name}"
+        # List all files in the repo
+        self.all_files = list_repo_files(repo_id, repo_type='model')
+        
+    def downloadbergamotfiles(self):
+        print('input text type:', type(self.input_text), len(self.all_files), 'installed_pairs', self.installed_pairs, 'defaultlocalfolder', self.localfolder)
+        try:
+            for branch in self.branches:
+                    branch_files = [f for f in self.all_files if f.startswith(branch)] 
+                    fullmodel_files = [f for f in branch_files if f.startswith(self.model_name)]
+                    print('branch_files', len(branch_files), 'fullmodel_files', fullmodel_files)
+                    model_files = [f.split(f'{self.model_name}/')[1] for f in fullmodel_files]
+                    print('branch_files', len(branch_files), 'model_files', model_files)
+                    for file_path in model_files:
+                        if localfolder not in installed_pairs:
+                            # local_files_only (bool, optional, defaults to False) — If True, avoid downloading the file and return the path to the local cached file if it exists.
+                            # dry_run (bool, optional, defaults to False) — If True, perform a dry run without actually downloading the file. Returns a DryRunFileInfo object containing information about what would be downloaded.
+                            local_path = hf_hub_download(repo_id=self.repo_id, subfolder=self.model_name, filename=file_path, local_dir=self.subfolder)
+                            print(f"Downloaded to: {local_path}") # Downloaded to: deen/base/deen/config.yml
+                ## Check if model/localfolder in repo files, add to set if exists
+                modelcheck = [i for i in all_files if self.model_name in i]
+                print('Modelcheck', modelcheck)
+                if modelcheck:
+                    print('Add to set after modelcheck', modelcheck)
+                    self.installed_pairs.add(self.localfolder)
+                try:
+                    dry_run = hf_hub_download(repo_id=self.repo_id, subfolder=self.model_name, filename='config.yml', local_dir=self.subfolder)
+                    print('installed_pairs', self.installed_pairs, 'localfolder', self.localfolder, 'dry_run', type(dry_run), dry_run)
+                    if isinstance(dry_run, str):
+                        print('Add to set after dryrun', dry_run)
+                        self.installed_pairs.add(localfolder)
+                except Exception as dryrunerror:
+                    print('installed_pairs', self.installed_pairs, 'localfolder', self.localfolder, 'dry_runerror', dryrunerror)
+        except Exception as downloaderror:
+        response, message_text = str(downloaderror), f"Error downloading {self.model_name}: {downloaderror}."
+        print(downloaderror)
+    def translate(self):
+    try:
+        model = service.modelFromConfigPath(f"{self.localfolder}/config.yml")
+        # model = service.modelFromConfig(localfolder)
+        options = bergamot.ResponseOptions(alignment=False, sentenceMappings=False, qualityScores=False, HTML=False)
+        rawresponse = service.translate(model, bergamot.VectorString(self.input_text), options)
+        response: list|str = [r.target.text for r in rawresponse] if len(rawresponse) > 1 else next(iter(rawresponse)).target.text
+        print(type(self.input_text), len(self.input_text), len(rawresponse), type(response), response)
+        # response = [r.target.text for r in model_response][0] if isinstance(response, bergamot._bergamot.VectorResponse) else next(iter(response)).target.text
+        # response is of type bergamot._bergamot.VectorResponse, an iterable of bergamot._bergamot.Response
+        message_text = f"Translated from {self.sl} to {self.tl} with {self.model_name}."
+    except Exception as translateerror:
+        response, message_text = str(translateerror), f"Error translating from {self.sl} to {self.tl} with {self.model_name}: {translateerror}."
+        print(translateerror)
+    return {"input": self.input_text, "translated_text": response, "message_text": message_text}
 
 # https://tiberiucristianleon-fastapimt.hf.space/bergamot?input_text=das%20ist%20keine%20gute%20Frau&input_text=das%20ist%20eine%20gute%20Nachricht&sl=de&tl=en&model=bergamot
 @app.get("/bergamot", operation_id="get_bergamot", description="Translate text with Bergamot", tags=["bergamot"], summary="Translate text with Bergamot")
@@ -101,17 +162,42 @@ def bergamot(input_text: list[str] = Query(description="Input string or list of
             translated_text(str): The input text translated into the selected target language
             message_text(str):  A descriptive message summarizing the translation process. Example: "Translated from English to German with base/ende."
     
+    Example:
+        >>> bergamot("Hello world", "en", "de", "base/ende")
+        {"input_text": "Hello world", "translated_text": "Hallo Welt", "message_text": "Translated from English to German with base/ende."}
+    """
+    try:
+        bergamotinstance = Bergamot(input_text, sl, tl, model_name)
+        bergamotinstance.downloadbergamotfiles()
+        return bergamotinstance.translate()
+
+# https://tiberiucristianleon-fastapimt.hf.space/bergamot?input_text=das%20ist%20keine%20gute%20Frau&input_text=das%20ist%20eine%20gute%20Nachricht&sl=de&tl=en&model=bergamot
+@app.get("/bergamots", operation_id="get_bergamot", description="Translate text with Bergamot", tags=["bergamot"], summary="Translate text with Bergamot")
+def bergamots(input_text: list[str] = Query(description="Input string or list of strings"), sl: str = 'de', tl: str = 'en', model_name: Optional[str] = 'base/deen'):
+    """
+    Translates the input text from the source language to the target language using a specified model.
+    Parameters:
+        input_text (str | list[str]): The source text to be translated, can be either a string or a list of strings
+        sl (str): The source language of the input text
+        tl (str): The target language into which the input text is translated
+        model_name (str): The selected translation model name
+    Returns:
+        dict:
+            input_text(str): The input text in the source language
+            translated_text(str): The input text translated into the selected target language
+            message_text(str):  A descriptive message summarizing the translation process. Example: "Translated from English to German with base/ende."
+    
     Example:
         >>> bergamot("Hello world", "en", "de", "base/ende")
         {"input_text": "Hello world", "translated_text": "Hallo Welt", "message_text": "Translated from English to German with base/ende."}
     """
     try:
         import bergamot
+        repo_id="TiberiuCristianLeon/Bergamot"
+        branches = ['base', 'base-memory', 'tiny']
         # input_text = [input_text] if isinstance(input_text, str) else input_text           
         config = bergamot.ServiceConfig(numWorkers=4)
         service = bergamot.Service(config)
-        repo_id="TiberiuCristianLeon/Bergamot"
-        branches = ['base', 'base-memory', 'tiny']
         subfolder = f"{sl}{tl}"
         localfolder = f"{subfolder}/{model_name}"
         # List all files in the repo