Update app.py

app.py

@@ -16,7 +16,7 @@ def init_model():
     model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0",
                                                  torch_dtype=torch.bfloat16, trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=30.)
     return model, tokenizer, streamer
 
 
@@ -30,14 +30,17 @@ def process(message, history):
     generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048, do_sample=True,
                              top_k=20, top_p=0.84, temperature=1.0, repetition_penalty=1.15, eos_token_id=2,
                              bos_token_id=1, pad_token_id=0)
-
-    t = Thread(target=model.generate, kwargs=generation_kwargs)
-    t.start()
-    response = ""
-    for text in streamer:
-        response += text
-        yield response
-    print('log:', response)
+    try:
+        t = Thread(target=model.generate, kwargs=generation_kwargs)
+        t.start()
+        response = ""
+        for text in streamer:
+            response += text
+            yield response
+        print('-log:', response)
+    except Exception as e:
+        print('-error:', str(e))
+        return "Error: 遇到错误，请开启新的会话重新尝试~"
 
 
 if __name__ == '__main__':