Add quick try Colab Notebook

README.md

@@ -230,6 +230,13 @@ chinese_llama_lora_7b/
 
 具体内容请参考本项目 >>> [📚 GitHub Wiki](https://github.com/ymcui/Chinese-LLaMA-Alpaca/wiki/训练细节)
 
+## 在colab快速体验
+
+| colab link                                                   | Descriptions                                                 |
+| ------------------------------------------------------------ | ------------------------------------------------------------ |
+| [![Open In Colab](https://camo.githubusercontent.com/84f0493939e0c4de4e6dbe113251b4bfb5353e57134ffd9fcab6b8714514d4d1/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/drive/1mG8CwGub-hcoeo79FjfoWtxqOuu3P_y1?usp=sharing) | 快速体验 |
+
+
 ## FAQ
 
 FAQ中给出了常见问题的解答，请在提Issue前务必先查看FAQ。

README_EN.md

@@ -230,6 +230,11 @@ We have open-sourced the scripts for pre-training and instruction finetuning (SF
 
 Please refer to our  >>> [📚GitHub Wiki](https://github.com/ymcui/Chinese-LLaMA-Alpaca/wiki/Training-Details).
 
+## Try on colab
+
+| colab link                                                   | Descriptions                                                 |
+| ------------------------------------------------------------ | ------------------------------------------------------------ |
+| [![Open In Colab](https://camo.githubusercontent.com/84f0493939e0c4de4e6dbe113251b4bfb5353e57134ffd9fcab6b8714514d4d1/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667)](https://colab.research.google.com/drive/1mG8CwGub-hcoeo79FjfoWtxqOuu3P_y1?usp=sharing) | quick try of chinese llama |
 
 ## FAQ
 

scripts/gradio_demo.py

@@ -2,15 +2,18 @@
 import gradio as gr
 import argparse
 import os
-import mdtex2html
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--base_model', default=None, type=str, required=True)
 parser.add_argument('--lora_model', default=None, type=str,help="If None, perform inference on the base model")
 parser.add_argument('--tokenizer_path',default=None,type=str)
 parser.add_argument('--gpus', default="0", type=str)
+parser.add_argument('--share', default=True, help='share gradio domain name')
+parser.add_argument('--load_in_8bit',action='store_true', help='use 8 bit model')
 parser.add_argument('--only_cpu',action='store_true',help='only use CPU for inference')
 args = parser.parse_args()
+share = args.share
+load_in_8bit = args.load_in_8bit
 if args.only_cpu is True:
     args.gpus = ""
 os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
@@ -19,19 +22,6 @@
 from transformers import LlamaForCausalLM, LlamaTokenizer, GenerationConfig
 from peft import PeftModel
 
-def postprocess(self, y):
-    if y is None:
-        return []
-    for i, (message, response) in enumerate(y):
-        y[i] = (
-            None if message is None else mdtex2html.convert((message)),
-            None if response is None else mdtex2html.convert(response),
-        )
-    return y
-
-
-gr.Chatbot.postprocess = postprocess
-
 generation_config = dict(
     temperature=0.2,
     top_k=40,
@@ -54,7 +44,7 @@ def postprocess(self, y):
 
 base_model = LlamaForCausalLM.from_pretrained(
     args.base_model, 
-    load_in_8bit=False,
+    load_in_8bit=load_in_8bit,
     torch_dtype=load_type,
     low_cpu_mem_usage=True,
     device_map='auto',
@@ -173,4 +163,4 @@ def predict(
     submitBtn.click(reset_user_input, [], [user_input])
 
     emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
-demo.queue().launch(share=False, inbrowser=True, server_name = '0.0.0.0', server_port=19324)
+demo.queue().launch(share=share, inbrowser=True, server_name = '0.0.0.0', server_port=19324)