From 3cc035634d75e78940f9c49f9ba6961fa67db397 Mon Sep 17 00:00:00 2001
From: Letong Han <106566639+letonghan@users.noreply.github.com>
Date: Fri, 26 Apr 2024 21:56:48 +0800
Subject: [PATCH] [Code Translation] Modify prompt and streaming format (#99)

Signed-off-by: letonghan <letong.han@intel.com>
---
 CodeTrans/README.md                                 | 2 +-
 CodeTrans/langchain/docker/build_docker.sh          | 2 +-
 CodeTrans/langchain/docker/codetrans-app/prompts.py | 4 ++++
 CodeTrans/langchain/docker/codetrans-app/server.py  | 7 ++++---
 4 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/CodeTrans/README.md b/CodeTrans/README.md
index 667ef8dc1..90b6197fa 100644
--- a/CodeTrans/README.md
+++ b/CodeTrans/README.md
@@ -23,7 +23,7 @@ bash launch_tgi_service.sh
 ```sh
 cd langchain/docker
 bash build_docker.sh
-docker run -it --name code_trans_server --net=host --ipc=host -e TGI_ENDPOINT=${TGI ENDPOINT} -e HUGGINGFACEHUB_API_TOKEN=${HUGGINGFACE_API_TOKEN} -e SERVER_PORT=8000 -e http_proxy=${http_proxy} -e https_proxy=${https_proxy} code_tranlation:latest bash
+docker run -it --name code_trans_server --net=host --ipc=host -e TGI_ENDPOINT=${TGI ENDPOINT} -e HUGGINGFACEHUB_API_TOKEN=${HUGGINGFACE_API_TOKEN} -e SERVER_PORT=8000 -e http_proxy=${http_proxy} -e https_proxy=${https_proxy} intel/gen-ai-examples:code-translation bash
 ```
 
 Here is the explanation of some of the above parameters:
diff --git a/CodeTrans/langchain/docker/build_docker.sh b/CodeTrans/langchain/docker/build_docker.sh
index 7f943df0e..1c058cbd3 100644
--- a/CodeTrans/langchain/docker/build_docker.sh
+++ b/CodeTrans/langchain/docker/build_docker.sh
@@ -14,4 +14,4 @@
 # limitations under the License.
 
 
-docker build . -t code_tranlation:latest --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy
+docker build . -t intel/gen-ai-examples:code-translation --build-arg https_proxy=$https_proxy --build-arg http_proxy=$http_proxy
diff --git a/CodeTrans/langchain/docker/codetrans-app/prompts.py b/CodeTrans/langchain/docker/codetrans-app/prompts.py
index 3442fe261..d3092f63c 100644
--- a/CodeTrans/langchain/docker/codetrans-app/prompts.py
+++ b/CodeTrans/langchain/docker/codetrans-app/prompts.py
@@ -18,8 +18,12 @@ prompt_template = """
     ### System: Please translate the following {language_from} codes into {language_to} codes.
 
     ### Original codes:
+    '''{language_from}
+
     {source_code}
 
+    '''
+
     ### Translated codes:
 """
 codetrans_prompt_template = PromptTemplate.from_template(prompt_template)
diff --git a/CodeTrans/langchain/docker/codetrans-app/server.py b/CodeTrans/langchain/docker/codetrans-app/server.py
index dd627fcdf..6e052bc3f 100644
--- a/CodeTrans/langchain/docker/codetrans-app/server.py
+++ b/CodeTrans/langchain/docker/codetrans-app/server.py
@@ -78,9 +78,10 @@ class CodeTranslationAPIRouter(APIRouter):
         print(f"[codetrans - stream] prompt:{prompt}")
 
         async def stream_generator():
-            async for chunk in self.llm.astream_log(prompt):
-                print(f"[codetrans - stream] data: {chunk}")
-                yield f"data: {chunk}\n\n"
+            for chunk in self.llm.stream(prompt):
+                chunk_repr = repr(chunk.encode("utf-8"))
+                print(f"[codetrans - stream] data: {chunk_repr}")
+                yield f"data: {chunk_repr}\n\n"
             yield "data: [DONE]\n\n"
 
         return StreamingResponse(stream_generator(), media_type="text/event-stream")