Add multilingual support

app.py

@@ -4,28 +4,65 @@ import gradio as gr
 import requests
 
 API_TOKEN = os.environ['API_TOKEN']
+G_TRANS_API_TOKEN = os.environ['G_TRANS_API_TOKEN']
 
-state = []
-
+API_URL = 'https://api-inference.huggingface.co/models/{}'
+G_TRANS_API = 'https://translation.googleapis.com/language/translate/v2'
 headers = {'Authorization': f'Bearer {API_TOKEN}'}
 
-def query(payload, model_id):
-    API_URL = f'https://api-inference.huggingface.co/models/{model_id}'
-    response = requests.post(API_URL, headers=headers, json=payload)
+def detect_lang(message):
+    response = requests.get(G_TRANS_API+'/detect', params={'key': G_TRANS_API_TOKEN, 'q': message})
+    return response.json()
+
+def translate_src_to_en(message, src_lang):
+    response = requests.get(G_TRANS_API, params={'key': G_TRANS_API_TOKEN, 'source': src_lang, 'target': 'en', 'q': message})
+    return response.json()
+
+def translate_en_to_src(message, src_lang):
+    response = requests.get(G_TRANS_API, params={'key': G_TRANS_API_TOKEN, 'source': 'en', 'target': src_lang, 'q': message})
+    return response.json()
+
+def query_model(model_id, payload):
+    response = requests.post(API_URL.format(model_id), headers=headers, json=payload)
     return response.json()
 
-def chat(message):
-    response = query({
-        'inputs': message,
+def parse_model_response(response):
+    return response[0]['generated_text']
+
+def parse_model_error(response):
+    return f'{response["error"]}. Please wait about {int(response["estimated_time"])} seconds.'
+
+def parse_translation_response(response):
+    return response['data']['translations'][0]['translatedText']
+
+def query_model(model_id, payload):
+    response = requests.post(API_URL.format(model_id), headers=headers, json=payload)
+    return response.json()
+
+state = []
+
+def chat(message, multi):
+    message_en = message
+    if multi:
+        response = detect_lang(message)
+        lang = response['data']['detections'][0][0]['language'][:2]
+        if lang != 'en':
+            response = translate_src_to_en(message, lang)
+            message_en = parse_translation_response(response)
+    response = query_model('IssakaAI/health-chatbot', {
+        'inputs': message_en,
         'parameters': {
             'max_length': 500,
         }
-    }, 'IssakaAI/health-chatbot')
+    })
     reply = ''
     if isinstance(response, list):
-        reply = response[0]['generated_text'][len(message) + 1:]
+        reply = parse_model_response(response)[len(message_en) + 1:]
+        if multi and lang != 'en':
+            response = translate_en_to_src(reply, lang)
+            reply = parse_translation_response(response)
     elif isinstance(response, dict):
-        reply = f'{response["error"]}. Please wait about {int(response["estimated_time"])} seconds.'
+        reply = parse_model_error(response)
     state.append((message, reply))
     return gr.Textbox.update(value=''), state
 
@@ -39,8 +76,9 @@ with gr.Blocks() as blk:
         chatbot = gr.Chatbot()
         with gr.Box():
             message = gr.Textbox(value='What is the menstrual cycle?', lines=10)
+            multi = gr.Checkbox(False, label='Multilingual chatbot')
             send = gr.Button('Send', variant='primary')
             clear = gr.Button('Clear history', variant='secondary')
-    send.click(fn=chat, inputs=message, outputs=[message, chatbot])
+    send.click(fn=chat, inputs=[message, multi], outputs=[message, chatbot])
     clear.click(fn=clear_message, inputs=[], outputs=chatbot)
 blk.launch(debug=True)