Update app.py

app.py

@@ -7,7 +7,7 @@ from hypercorn.asyncio import serve
 from hypercorn.config import Config
 import os
 os.environ['CURL_CA_BUNDLE'] = ''
-from googletranslate import translate
+#from googletranslate import translate
 import json
 import random
 import re
@@ -159,7 +159,7 @@ def random_spanish_pair2():
 
     
 
-
+"""
 @app.route('/translate', methods=['POST'])
 def dotranslate():
     data = request.get_json()
@@ -177,7 +177,48 @@ def dotranslate():
         return jsonify({'translation': translation}), 200
     else:
         return jsonify({'error': 'No text provided'}), 400
-        
+"""
+
+
+from transformers import M2M100ForConditionalGeneration
+from tokenization_small100 import SMALL100Tokenizer
+
+model_name = "alirezamsh/small100"
+model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+tokenizer = SMALL100Tokenizer.from_pretrained(model_name)
+
+@app.route('/translate', methods=['POST'])
+def dotranslate():
+    data = request.get_json()
+    txt = data.get('txt')
+    src = data.get('src', 'en')
+    dest = data.get('dest', 'es')
+
+    if txt:
+        cache_key = f"{txt}_{src}_{dest}"
+        translation = cache.get(cache_key)
+        if translation is None:
+            # Set the source and target languages
+            tokenizer.src_lang = src
+            tokenizer.tgt_lang = dest
+
+            # Tokenize the input text
+            encoded = tokenizer(txt, return_tensors="pt")
+
+            # Generate translation
+            generated_tokens = model.generate(
+                **encoded,
+                forced_bos_token_id=tokenizer.get_lang_id(dest)
+            )
+
+            # Decode the generated tokens
+            translation = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
+
+            # Cache the translation
+            cache.set(cache_key, translation)
+        return jsonify({'translation': translation}), 200
+    else:
+        return jsonify({'error': 'No text provided'}), 400
 
 if __name__ == "__main__":
    config = Config()