updated with 8 bit version

app.py

@@ -1,25 +1,34 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from modelscope.pipelines import pipeline
-from modelscope.utils.constant import Tasks
+tokenizer = AutoTokenizer.from_pretrained("baichuan-inc/baichuan-7B", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("baichuan-inc/baichuan-7B", trust_remote_code=True, device_map="auto", load_in_8bit=True)
 
-text_generation_zh = pipeline(task=Tasks.text_generation, model='baichuan-inc/baichuan-7B',model_revision='v1.0.2')
-text_generation_zh._model_prepare = True
+#from transformers import 
+#from modelscope.pipelines import pipeline
+#from modelscope.utils.constant import Tasks
+#text_generation_zh = pipeline(task=Tasks.text_generation, model='baichuan-inc/baichuan-7B',model_revision='v1.0.2')
+#text_generation_zh._model_prepare = True
 
 def text_generation(input_text):
     if input_text == "":
         return ""
-    return text_generation_zh(input_text)['text']
+    inputs = tokenizer(input_text, return_tensors='pt').to('cuda:0') #"run->walk\nwrite->"
+    pred = model.generate(**inputs, max_new_tokens=64, repetition_penalty=1.1)
+    return tokenizer.decode(pred.cpu()[0], skip_special_tokens=True) #text_generation_zh(input_text)['text']
 
 with gr.Blocks() as demo:
     gr.Markdown(
     """
-    # 百川-7B 模型体验
+    # [Chinese] 百川-7B 模型体验
     输入文本以查看生成结果。
+
+    # [English] Baichuan-7B model experience
+     Enter text to see build results.
     """)
-    inp = gr.Textbox(label="输入prompt")
-    submit = gr.Button("提交")
-    out = gr.Textbox(label="续写结果")
+    inp = gr.Textbox(label="输入Prompt / Input Prompt")
+    submit = gr.Button("提交/Submit")
+    out = gr.Textbox(label="续写结果/Generated Text")
     submit.click(text_generation, inp, out)
 
     gr.Markdown("## Text Examples")