oooooooooo

Dockerfile

@@ -1,3 +1,4 @@
+
 FROM python:3.9
 
 WORKDIR /app
@@ -8,7 +9,11 @@ RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY . .
 
-EXPOSE 5000
+# Set environment variable for transformers cache
+ENV TRANSFORMERS_CACHE=/app/cache
 
-CMD ["python","app.py"]
+RUN mkdir -p /app/cache
+
+EXPOSE 5000
 
+CMD ["python", "app.py"]

app.py

@@ -1,88 +1,3 @@
-# import os
-# import torch
-# import gradio s gr
-# import time
-# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-# from flores200_codes import flores_codes
-
-
-# def load_models():
-#     # build model and tokenizer
-#     model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
-#                   #'nllb-1.3B': 'facebook/nllb-200-1.3B',
-#                   #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
-#                   #'nllb-3.3B': 'facebook/nllb-200-3.3B',
-#                   }
-
-#     model_dict = {}
-
-#     for call_name, real_name in model_name_dict.items():
-#         print('\tLoading model: %s' % call_name)
-#         model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
-#         tokenizer = AutoTokenizer.from_pretrained(real_name)
-#         model_dict[call_name+'_model'] = model
-#         model_dict[call_name+'_tokenizer'] = tokenizer
-
-#     return model_dict
-
-
-# def translation(source, target, text):
-#     if len(model_dict) == 2:
-#         model_name = 'nllb-distilled-600M'
-
-#     start_time = time.time()
-#     source = flores_codes[source]
-#     target = flores_codes[target]
-
-#     model = model_dict[model_name + '_model']
-#     tokenizer = model_dict[model_name + '_tokenizer']
-
-#     translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
-#     output = translator(text, max_length=400)
-
-#     end_time = time.time()
-
-#     output = output[0]['translation_text']
-#     result = {'inference_time': end_time - start_time,
-#               'source': source,
-#               'target': target,
-#               'result': output}
-#     return result
-
-
-# if __name__ == '__main__':
-#     print('\tinit models')
-
-#     global model_dict
-
-#     model_dict = load_models()
-    
-#     # define gradio demo
-#     lang_codes = list(flores_codes.keys())
-#     #inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
-#     inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
-#               gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'),
-#               gr.inputs.Textbox(lines=5, label="Input text"),
-#               ]
-
-#     outputs = gr.outputs.JSON()
-
-#     title = "NLLB distilled 600M demo"
-
-#     demo_status = "Demo is running on CPU"
-#     description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
-#     examples = [
-#     ['English', 'Korean', 'Hi. nice to meet you']
-#     ]
-
-#     gr.Interface(translation,
-#                  inputs,
-#                  outputs,
-#                  title=title,
-#                  description=description,
-#                  ).launch() 
-
-
 import os
 import time
 from flask import Flask, request, jsonify
@@ -91,56 +6,66 @@ from flores200_codes import flores_codes
 
 app = Flask(__name__)
 
+
 def load_models():
-    model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M'}
+    model_name_dict = {"nllb-distilled-600M": "facebook/nllb-200-distilled-600M"}
     model_dict = {}
 
     for call_name, real_name in model_name_dict.items():
-        print(f'\tLoading model: {call_name}')
+        print(f"\tLoading model: {call_name}")
         model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
         tokenizer = AutoTokenizer.from_pretrained(real_name)
-        model_dict[call_name + '_model'] = model
-        model_dict[call_name + '_tokenizer'] = tokenizer
+        model_dict[call_name + "_model"] = model
+        model_dict[call_name + "_tokenizer"] = tokenizer
 
     return model_dict
 
+
 global model_dict
 model_dict = load_models()
 
-@app.route('/api/translate', methods=['POST'])
+
+@app.route("/api/translate", methods=["POST"])
 def translate_text():
     data = request.json
-    source_lang = data.get('source')
-    target_lang = data.get('target')
-    input_text = data.get('text')
-    
+    source_lang = data.get("source")
+    target_lang = data.get("target")
+    input_text = data.get("text")
+
     if not source_lang or not target_lang or not input_text:
         return jsonify({"error": "source, target, and text fields are required"}), 400
-    
-    model_name = 'nllb-distilled-600M'
+
+    model_name = "nllb-distilled-600M"
     start_time = time.time()
     source = flores_codes.get(source_lang)
     target = flores_codes.get(target_lang)
-    
+
     if not source or not target:
         return jsonify({"error": "Invalid source or target language code"}), 400
 
-    model = model_dict[model_name + '_model']
-    tokenizer = model_dict[model_name + '_tokenizer']
-    
-    translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
+    model = model_dict[model_name + "_model"]
+    tokenizer = model_dict[model_name + "_tokenizer"]
+
+    translator = pipeline(
+        "translation",
+        model=model,
+        tokenizer=tokenizer,
+        src_lang=source,
+        tgt_lang=target,
+    )
     output = translator(input_text, max_length=400)
-    
+
     end_time = time.time()
-    output_text = output[0]['translation_text']
-    
+    output_text = output[0]["translation_text"]
+
     result = {
-        'inference_time': end_time - start_time,
-        'source': source_lang,
-        'target': target_lang,
-        'result': output_text
+        "inference_time": end_time - start_time,
+        "source": source_lang,
+        "target": target_lang,
+        "result": output_text,
     }
     return jsonify(result)
 
-if __name__ == '__main__':
-    app.run(debug=True) 
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000, debug=True)