ClueAI commited on
Commit
f11ab26
1 Parent(s): 0a57310

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -3,9 +3,10 @@ import gradio as gr
3
  import clueai
4
  from transformers import T5Tokenizer, T5ForConditionalGeneration
5
  tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v2")
6
- model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v2")
7
  # 使用
8
- device='cpu'
 
9
 
10
  def preprocess(text):
11
  text = text.replace("\n", "\\n").replace("\t", "\\t")
 
3
  import clueai
4
  from transformers import T5Tokenizer, T5ForConditionalGeneration
5
  tokenizer = T5Tokenizer.from_pretrained("ClueAI/ChatYuan-large-v2")
6
+ model = T5ForConditionalGeneration.from_pretrained("ClueAI/ChatYuan-large-v2").half()
7
  # 使用
8
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
+
10
 
11
  def preprocess(text):
12
  text = text.replace("\n", "\\n").replace("\t", "\\t")