developer3000 commited on
Commit
6ac2a7b
·
1 Parent(s): c7d54c9

Add application file

Browse files
Files changed (1) hide show
  1. app.py +1 -2
app.py CHANGED
@@ -5,7 +5,6 @@ from threading import Thread
5
 
6
  tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
7
  model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
8
- model = model.to('cuda:0')
9
 
10
  class StopOnTokens(StoppingCriteria):
11
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
@@ -22,7 +21,7 @@ def predict(message, history):
22
  messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
23
  for item in history_transformer_format])
24
 
25
- model_inputs = tokenizer([messages], return_tensors="pt").to("cuda")
26
  streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
27
  generate_kwargs = dict(
28
  model_inputs,
 
5
 
6
  tokenizer = AutoTokenizer.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1")
7
  model = AutoModelForCausalLM.from_pretrained("togethercomputer/RedPajama-INCITE-Chat-3B-v1", torch_dtype=torch.float16)
 
8
 
9
  class StopOnTokens(StoppingCriteria):
10
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
 
21
  messages = "".join(["".join(["\n<human>:"+item[0], "\n<bot>:"+item[1]])
22
  for item in history_transformer_format])
23
 
24
+ model_inputs = tokenizer([messages], return_tensors="pt")
25
  streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
26
  generate_kwargs = dict(
27
  model_inputs,