Fix chatglm (#13)

* Update chatglm2.py * Update app.py * parse_codeblock
ypwhs · Apr 1, 2023 · a9dbbe6 · a9dbbe6
1 parent 20d71c6
commit a9dbbe6
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/app.py b/app.py
@@ -14,7 +14,7 @@
 # model_name = 'BelleGroup/BELLE-LLAMA-7B-2M-gptq'
 
 if 'chatglm' in model_name.lower():
-    from predictors.chatglm import ChatGLM
+    from predictors.chatglm2 import ChatGLM
     predictor = ChatGLM(model_name)
 elif 'gptq' in model_name.lower():
     from predictors.llama_gptq import LLaMaGPTQ
@@ -26,7 +26,7 @@
     from predictors.debug import Debug
     predictor = Debug(model_name)
 else:
-    from predictors.chatglm import ChatGLM
+    from predictors.chatglm2 import ChatGLM
     predictor = ChatGLM(model_name)
 
 

diff --git a/predictors/base.py b/predictors/base.py
@@ -1,6 +1,20 @@
 from abc import ABC, abstractmethod
 
 
+def parse_codeblock(text):
+    lines = text.split("\n")
+    for i, line in enumerate(lines):
+        if "```" in line:
+            if line != "```":
+                lines[i] = f'<pre><code class="{lines[i][3:]}">'
+            else:
+                lines[i] = '</code></pre>'
+        else:
+            if i > 0:
+                lines[i] = "<br/>" + line.replace("<", "&lt;").replace(">", "&gt;")
+    return "".join(lines)
+
+
 class BasePredictor(ABC):
 
     @abstractmethod

diff --git a/predictors/chatglm.py → predictors/chatglm2.py b/predictors/chatglm.py → predictors/chatglm2.py
@@ -5,7 +5,8 @@
 from transformers import AutoModel, AutoTokenizer
 from transformers import LogitsProcessor, LogitsProcessorList
 
-from predictors.base import BasePredictor
+from predictors.base import BasePredictor, parse_codeblock
+from chatglm.modeling_chatglm import ChatGLMForConditionalGeneration
 
 
 class InvalidScoreLogitsProcessor(LogitsProcessor):
@@ -27,7 +28,7 @@ def __init__(self, model_name):
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_name, trust_remote_code=True, resume_download=True)
         if 'int4' not in model_name:
-            model = AutoModel.from_pretrained(
+            model = ChatGLMForConditionalGeneration.from_pretrained(
                 model_name,
                 trust_remote_code=True,
                 resume_download=True,
@@ -36,7 +37,7 @@ def __init__(self, model_name):
                 device_map={'': self.device}
             )
         else:
-            model = AutoModel.from_pretrained(
+            model = ChatGLMForConditionalGeneration.from_pretrained(
                 model_name,
                 trust_remote_code=True,
                 resume_download=True
@@ -105,4 +106,4 @@ def stream_chat_continue(self,
             outputs = outputs.tolist()[0][input_length:]
             response = tokenizer.decode(outputs)
             response = model.process_response(response)
-            yield response
+            yield parse_codeblock(response)