Update app.py

app.py

@@ -1,64 +1,101 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+LANGUAGES = ["en", "de", "es", "fr", "it", "nl", "sv", "pt"]
+DOMAINS = {
+    "Asset management": "am",
+    "Annual report": "ar",
+    "Corporate action": "corporateAction",
+    "Equity research": "equi",
+    "Fund fact sheet": "ffs",
+    "Kiid": "kiid",
+    "Life insurance": "lifeInsurance",
+    "Regulatory": "regulatory",
+    "General": "general",
+}
 
+# Helper functions
+def language_token(lang):
+    return f"<lang_{lang}>"
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+def domain_token(dom):
+    return f"<dom_{dom}>"
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+def format_input(src, tgt_lang, src_lang, domain):
+    assert tgt_lang in LANGUAGES
+    tgt_lang_token = language_token(tgt_lang)
+    # Prefix the input with <eos>
+    base_input = f"<eos>{src}</src>{tgt_lang_token}"
+    if src_lang:
+        assert src_lang in LANGUAGES
+        src_lang_token = language_token(src_lang)
+        base_input = f"{base_input}{src_lang_token}"
+    if domain:
+        domain = DOMAINS.get(domain, "general")
+        dom_token = domain_token(domain)
+        base_input = f"{base_input}{dom_token}"
+    return base_input
 
-    messages.append({"role": "user", "content": message})
+# Initialize model and tokenizer globally to avoid reloading
+model_id = "LinguaCustodia/multilingual-multidomain-fin-mt-70M"
+auth_token = os.environ.get("TOKEN") or True
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=auth_token)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
-    response = ""
+def translate(text, source_lang, target_lang, domain):
+    if not text:
+        return ""
+    
+    src_lang_code = language_map.get(source_lang)
+    tgt_lang_code = language_map.get(target_lang)
+    
+    formatted_sentence = format_input(text, tgt_lang_code, src_lang_code, domain)
+    inputs = tokenizer(formatted_sentence, return_tensors="pt", return_token_type_ids=False)
+    
+    outputs = model.generate(**inputs, max_new_tokens=256)
+    
+    input_size = inputs["input_ids"].size(1)
+    translated_sentence = tokenizer.decode(
+        outputs[0, input_size:], skip_special_tokens=True
+    )
+    
+    return translated_sentence
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+language_map = {
+    "English": "en", 
+    "German": "de", 
+    "Spanish": "es", 
+    "French": "fr", 
+    "Italian": "it", 
+    "Dutch": "nl", 
+    "Swedish": "sv", 
+    "Portuguese": "pt"
+}
 
-        response += token
-        yield response
+title = "🌐 Multilingual Multidomain Financial Translator 🌐"
+description = """<p><center>Specialized Translation for Financial Documents across 8 Languages and 9 Domains</center></p>"""
+article = """<p style='text-align: center'>Model: <a href='https://huggingface.co/LinguaCustodia/multilingual-multidomain-fin-mt-70M' target='_blank'>LinguaCustodia/multilingual-multidomain-fin-mt-70M</a></p>"""
 
+examples = [
+    ["Nous avons enregistré une croissance du chiffre d'affaires de 5,7% au troisième trimestre.", "French", "English", "Annual report"],
+    ["The funds under management increased by €2.3 billion during the fiscal year.", "English", "Spanish", "Asset management"],
+    ["Der Aufsichtsrat hat den Jahresabschluss geprüft und genehmigt.", "German", "French", "Regulatory"]
+]
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+demo = gr.Interface(
+    fn=translate,
+    title=title,
+    description=description,
+    article=article,
+    inputs=[
+        gr.Textbox(lines=5, placeholder="Enter text to translate (maximum 5 lines)", label="Input Text"),
+        gr.Dropdown(choices=list(language_map.keys()), value="French", label="Source Language"),
+        gr.Dropdown(choices=list(language_map.keys()), value="English", label="Target Language"),
+        gr.Dropdown(choices=list(DOMAINS.keys()), value="General", label="Financial Domain"),
     ],
+    outputs=gr.Textbox(label="Translation"),
+    examples=examples
 )
 
-
-if __name__ == "__main__":
-    demo.launch()
+demo.launch(enable_queue=True)