trans_api

app.py

@@ -12,7 +12,7 @@ from transformers import AutoModel, AutoTokenizer
 import gradio as gr
 import mdtex2html
 
-# from loguru import logger
+from loguru import logger
 
 model_name = "THUDM/chatglm2-6b"
 model_name = "THUDM/chatglm2-6b-int4"
@@ -96,6 +96,23 @@ def predict(input, chatbot, max_length, top_p, temperature, history, past_key_va
         yield chatbot, history, past_key_values
 
 
+def trans_api(input, max_length=4096, top_p=0.8, temperature=0.2):
+    try:
+        res = model.stream_chat(
+            tokenizer, 
+            input, 
+            history=[], 
+            past_key_values=None,
+            return_past_key_values=False,
+            max_length=max_length,
+            top_p=top_p,
+            temperature=temperature,
+        )
+        logger.debug(f"{res=}")
+    except Exception as exc:
+        logger.error(exc)
+        
+
 def reset_user_input():
     return gr.update(value='')
 
@@ -147,6 +164,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     emptyBtn.click(reset_state, outputs=[chatbot, history, past_key_values], show_progress=True)
 
+    with gr.Accordion("For Translation API", open=False):
+        input_text = gr.Text()
+        tr_btn = gr.Button("Go", variant="primary")
+    tr_btn.click(trans_api, [input_text, max_length, top_p, temperature], [], show_progress=True, api_name="tr")
+        
 # demo.queue().launch(share=False, inbrowser=True)
 # demo.queue().launch(share=True, inbrowser=True, debug=True)