Update
Browse files
model.py
CHANGED
@@ -34,7 +34,7 @@ def run(message: str,
|
|
34 |
top_p: float = 0.95,
|
35 |
top_k: int = 50) -> Iterator[str]:
|
36 |
prompt = get_prompt(message, chat_history, system_prompt)
|
37 |
-
inputs = tokenizer([prompt], return_tensors='pt').to(
|
38 |
|
39 |
streamer = TextIteratorStreamer(tokenizer,
|
40 |
timeout=10.,
|
|
|
34 |
top_p: float = 0.95,
|
35 |
top_k: int = 50) -> Iterator[str]:
|
36 |
prompt = get_prompt(message, chat_history, system_prompt)
|
37 |
+
inputs = tokenizer([prompt], return_tensors='pt').to('cuda')
|
38 |
|
39 |
streamer = TextIteratorStreamer(tokenizer,
|
40 |
timeout=10.,
|