Hamza-cpp commited on
Commit
b05fcfe
·
1 Parent(s): e69a4b4

oooooooooo

Browse files
Files changed (2) hide show
  1. Dockerfile +7 -2
  2. app.py +35 -110
Dockerfile CHANGED
@@ -1,3 +1,4 @@
 
1
  FROM python:3.9
2
 
3
  WORKDIR /app
@@ -8,7 +9,11 @@ RUN pip install --no-cache-dir --upgrade -r requirements.txt
8
 
9
  COPY . .
10
 
11
- EXPOSE 5000
 
12
 
13
- CMD ["python","app.py"]
 
 
14
 
 
 
1
+
2
  FROM python:3.9
3
 
4
  WORKDIR /app
 
9
 
10
  COPY . .
11
 
12
+ # Set environment variable for transformers cache
13
+ ENV TRANSFORMERS_CACHE=/app/cache
14
 
15
+ RUN mkdir -p /app/cache
16
+
17
+ EXPOSE 5000
18
 
19
+ CMD ["python", "app.py"]
app.py CHANGED
@@ -1,88 +1,3 @@
1
- # import os
2
- # import torch
3
- # import gradio s gr
4
- # import time
5
- # from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
6
- # from flores200_codes import flores_codes
7
-
8
-
9
- # def load_models():
10
- # # build model and tokenizer
11
- # model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M',
12
- # #'nllb-1.3B': 'facebook/nllb-200-1.3B',
13
- # #'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B',
14
- # #'nllb-3.3B': 'facebook/nllb-200-3.3B',
15
- # }
16
-
17
- # model_dict = {}
18
-
19
- # for call_name, real_name in model_name_dict.items():
20
- # print('\tLoading model: %s' % call_name)
21
- # model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
22
- # tokenizer = AutoTokenizer.from_pretrained(real_name)
23
- # model_dict[call_name+'_model'] = model
24
- # model_dict[call_name+'_tokenizer'] = tokenizer
25
-
26
- # return model_dict
27
-
28
-
29
- # def translation(source, target, text):
30
- # if len(model_dict) == 2:
31
- # model_name = 'nllb-distilled-600M'
32
-
33
- # start_time = time.time()
34
- # source = flores_codes[source]
35
- # target = flores_codes[target]
36
-
37
- # model = model_dict[model_name + '_model']
38
- # tokenizer = model_dict[model_name + '_tokenizer']
39
-
40
- # translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
41
- # output = translator(text, max_length=400)
42
-
43
- # end_time = time.time()
44
-
45
- # output = output[0]['translation_text']
46
- # result = {'inference_time': end_time - start_time,
47
- # 'source': source,
48
- # 'target': target,
49
- # 'result': output}
50
- # return result
51
-
52
-
53
- # if __name__ == '__main__':
54
- # print('\tinit models')
55
-
56
- # global model_dict
57
-
58
- # model_dict = load_models()
59
-
60
- # # define gradio demo
61
- # lang_codes = list(flores_codes.keys())
62
- # #inputs = [gr.inputs.Radio(['nllb-distilled-600M', 'nllb-1.3B', 'nllb-distilled-1.3B'], label='NLLB Model'),
63
- # inputs = [gr.inputs.Dropdown(lang_codes, default='English', label='Source'),
64
- # gr.inputs.Dropdown(lang_codes, default='Korean', label='Target'),
65
- # gr.inputs.Textbox(lines=5, label="Input text"),
66
- # ]
67
-
68
- # outputs = gr.outputs.JSON()
69
-
70
- # title = "NLLB distilled 600M demo"
71
-
72
- # demo_status = "Demo is running on CPU"
73
- # description = f"Details: https://github.com/facebookresearch/fairseq/tree/nllb. {demo_status}"
74
- # examples = [
75
- # ['English', 'Korean', 'Hi. nice to meet you']
76
- # ]
77
-
78
- # gr.Interface(translation,
79
- # inputs,
80
- # outputs,
81
- # title=title,
82
- # description=description,
83
- # ).launch()
84
-
85
-
86
  import os
87
  import time
88
  from flask import Flask, request, jsonify
@@ -91,56 +6,66 @@ from flores200_codes import flores_codes
91
 
92
  app = Flask(__name__)
93
 
 
94
  def load_models():
95
- model_name_dict = {'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M'}
96
  model_dict = {}
97
 
98
  for call_name, real_name in model_name_dict.items():
99
- print(f'\tLoading model: {call_name}')
100
  model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
101
  tokenizer = AutoTokenizer.from_pretrained(real_name)
102
- model_dict[call_name + '_model'] = model
103
- model_dict[call_name + '_tokenizer'] = tokenizer
104
 
105
  return model_dict
106
 
 
107
  global model_dict
108
  model_dict = load_models()
109
 
110
- @app.route('/api/translate', methods=['POST'])
 
111
  def translate_text():
112
  data = request.json
113
- source_lang = data.get('source')
114
- target_lang = data.get('target')
115
- input_text = data.get('text')
116
-
117
  if not source_lang or not target_lang or not input_text:
118
  return jsonify({"error": "source, target, and text fields are required"}), 400
119
-
120
- model_name = 'nllb-distilled-600M'
121
  start_time = time.time()
122
  source = flores_codes.get(source_lang)
123
  target = flores_codes.get(target_lang)
124
-
125
  if not source or not target:
126
  return jsonify({"error": "Invalid source or target language code"}), 400
127
 
128
- model = model_dict[model_name + '_model']
129
- tokenizer = model_dict[model_name + '_tokenizer']
130
-
131
- translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=source, tgt_lang=target)
 
 
 
 
 
 
132
  output = translator(input_text, max_length=400)
133
-
134
  end_time = time.time()
135
- output_text = output[0]['translation_text']
136
-
137
  result = {
138
- 'inference_time': end_time - start_time,
139
- 'source': source_lang,
140
- 'target': target_lang,
141
- 'result': output_text
142
  }
143
  return jsonify(result)
144
 
145
- if __name__ == '__main__':
146
- app.run(debug=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
2
  import time
3
  from flask import Flask, request, jsonify
 
6
 
7
  app = Flask(__name__)
8
 
9
+
10
  def load_models():
11
+ model_name_dict = {"nllb-distilled-600M": "facebook/nllb-200-distilled-600M"}
12
  model_dict = {}
13
 
14
  for call_name, real_name in model_name_dict.items():
15
+ print(f"\tLoading model: {call_name}")
16
  model = AutoModelForSeq2SeqLM.from_pretrained(real_name)
17
  tokenizer = AutoTokenizer.from_pretrained(real_name)
18
+ model_dict[call_name + "_model"] = model
19
+ model_dict[call_name + "_tokenizer"] = tokenizer
20
 
21
  return model_dict
22
 
23
+
24
  global model_dict
25
  model_dict = load_models()
26
 
27
+
28
+ @app.route("/api/translate", methods=["POST"])
29
  def translate_text():
30
  data = request.json
31
+ source_lang = data.get("source")
32
+ target_lang = data.get("target")
33
+ input_text = data.get("text")
34
+
35
  if not source_lang or not target_lang or not input_text:
36
  return jsonify({"error": "source, target, and text fields are required"}), 400
37
+
38
+ model_name = "nllb-distilled-600M"
39
  start_time = time.time()
40
  source = flores_codes.get(source_lang)
41
  target = flores_codes.get(target_lang)
42
+
43
  if not source or not target:
44
  return jsonify({"error": "Invalid source or target language code"}), 400
45
 
46
+ model = model_dict[model_name + "_model"]
47
+ tokenizer = model_dict[model_name + "_tokenizer"]
48
+
49
+ translator = pipeline(
50
+ "translation",
51
+ model=model,
52
+ tokenizer=tokenizer,
53
+ src_lang=source,
54
+ tgt_lang=target,
55
+ )
56
  output = translator(input_text, max_length=400)
57
+
58
  end_time = time.time()
59
+ output_text = output[0]["translation_text"]
60
+
61
  result = {
62
+ "inference_time": end_time - start_time,
63
+ "source": source_lang,
64
+ "target": target_lang,
65
+ "result": output_text,
66
  }
67
  return jsonify(result)
68
 
69
+
70
+ if __name__ == "__main__":
71
+ app.run(host="0.0.0.0", port=5000, debug=True)