Update app.py

app.py

@@ -2,13 +2,22 @@ import os
 import gradio as gr
 from langchain.llms import HuggingFaceHub
 
-model_repo = os.getenv('HF_MODEL_REPO')
-template = """[INST]<<SYS>>I want you to act as document language translator. You do translation {source} texts in document into then you return to me the translated document AND DO NOTHING ELSE.<</SYS>>
+model_repo = os.getenv('HF_MODEL_LLAMA_REPO')
+model_repo = os.getenv('HF_MODEL_STARCHAT_REPO')
+llamma_template = """[INST]<<SYS>>I want you to act as document language translator. You do translation {source} texts in document into then you return to me the translated document AND DO NOTHING ELSE.<</SYS>>
 Begin of the document:
 {query}
 End of the document[/INST]
 {target} translated document:
 
+"""
+starchat_template = """<|system|>I want you to act as document language translator. You do translation {source} texts in document into then you return to me the translated document AND DO NOTHING ELSE.<</SYS>>
+Begin of the document:
+{query}
+End of the document<|end|>
+<|assistant|>
+{target} translated document:
+
 """
 
 model_kwargs={
@@ -19,14 +28,22 @@ model_kwargs={
             "stop" : ["</s>","<|endoftext|>","<|end|>"],
             }
 
-llm = HuggingFaceHub(repo_id=model_repo, task="text-generation", model_kwargs=model_kwargs)
+llm1 = HuggingFaceHub(repo_id=model_repo, task="text-generation", model_kwargs=model_kwargs)
+llm2 = HuggingFaceHub(repo_id=model_repo, task="text-generation", model_kwargs=model_kwargs)
 
 def translation(source, target, text):
-    input_prompt = template.replace("{source}", source)
-    input_prompt = input_prompt.replace("{target}", target)
-    input_prompt = input_prompt.replace("{query}", text)
-    print(input_prompt)
-    response=llm(input_prompt)
+    response = text
+    try:
+        input_prompt = llamma_template.replace("{source}", source)
+        input_prompt = input_prompt.replace("{target}", target)
+        input_prompt = input_prompt.replace("{query}", text)
+        response=llm1(input_prompt)
+    except Exception as e:
+        print(f"ERROR: LLM show {e}")
+        input_prompt = starchat_template.replace("{source}", source)
+        input_prompt = input_prompt.replace("{target}", target)
+        input_prompt = input_prompt.replace("{query}", text)
+        response=llm2(input_prompt)
     return response
 
 gr.Interface(translation, inputs=["text","text","text"], outputs="text").launch()