Update app.py

app.py

@@ -1,8 +1,60 @@
 import gradio as gr
-import chatbot
-print(chatbot.__version__)
-model = chatbot.Bot()
-model.load("szzzzz/chatbot_bloom_560m",low_disk_usage=True)
+import torch
+import requests
+from transformers import BloomForCausalLM, BloomTokenizerFast
+import os 
+
+repo_id = 'szzzzz/chatbot_bloom_560m'
+os.mkdir('./chatbot')
+path = huggingface_hub.snapshot_download(
+                repo_id=repo_id, cache_dir='./chatbot',ignore_patterns = "*bin"
+            )
+url = huggingface_hub.file_download.hf_hub_url(repo_id, "pytorch_model.bin")
+tokenizer = BloomTokenizerFast.from_pretrained(path)
+state_dict = torch.load(
+    io.BytesIO(requests.get(url).content), map_location=torch.device("cpu")
+)
+model = BloomForCausalLM.from_pretrained(
+    pretrained_model_name_or_path=None,
+    state_dict=state_dict,
+    config=f"{path}/config.json",
+)
+max_length=1024
+
+
+def generate(inputs: str) -> str:
+    """generate content on inputs .
+
+    Args:
+        inputs (str):
+            example :'Human: 你好 .\n \nAssistant: '
+
+    Returns:
+        str:
+            bot response
+            example : '你好!我是你的ai助手!'
+
+    """
+    input_text = tokenizer.bos_token + inputs
+    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    _, input_len = input_ids.shape
+    if input_len >= max_length - 4:
+        res = "对话超过字数限制,请重新开始."
+        return res
+    pred_ids = model.generate(
+        input_ids,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        bos_token_id=tokenizer.bos_token_id,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.8,
+        max_new_tokens=max_length - input_len,
+        repetition_penalty=1.2,
+    )
+    pred = pred_ids[0][input_len:]
+    res = tokenizer.decode(pred, skip_special_tokens=True)
+    return res
 
 
 def add_text(history, text):
@@ -19,7 +71,7 @@ def bot(history):
         else:
             prompt = prompt + "\nAssistant: "
 
-    response = model.generate(prompt)
+    response = generate(prompt)
     history[-1][1] = response
     return history
 
@@ -32,7 +84,7 @@ def regenerate(history):
         else:
             prompt = prompt + "\nAssistant: "
 
-    response = model.generate(prompt)
+    response = generate(prompt)
     history[-1][1] = response
     return history