Akjava commited on
Commit
28a3da2
Β·
verified Β·
1 Parent(s): af67c6c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -60,7 +60,7 @@ def generate_text(messages):
60
  model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
61
  )
62
  streamer = TextStreamer(tokenizer, skip_prompt=True)
63
- text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ,streamer=streamer) #pipeline has not to(device)
64
  result = text_generator(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
65
  print(f"result={result}")
66
  generated_output = ""
 
60
  model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
61
  )
62
  streamer = TextStreamer(tokenizer, skip_prompt=True)
63
+ text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ,streamer=streamer,model_kwargs={"stream": True} ) #pipeline has not to(device)
64
  result = text_generator(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
65
  print(f"result={result}")
66
  generated_output = ""