Spaces:

LLM-DEMO
/

llama-2-medical-consultation

Paused

Ashishkr commited on Sep 4, 2023

Commit

1312c32

•

1 Parent(s): 18ec5f7

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -70,7 +70,7 @@ def run(message: str,
         top_p: float = 0.95,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
-    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
     streamer = TextIteratorStreamer(tokenizer,
                                     timeout=10.,

         top_p: float = 0.95,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
+    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to(device)
     streamer = TextIteratorStreamer(tokenizer,
                                     timeout=10.,