dar-tau commited on
Commit
7a53c83
·
verified ·
1 Parent(s): 6b3281f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -53,8 +53,7 @@ Assistant: girlfriend;mother;father;friend
53
  torch.set_grad_enabled(False)
54
  device = "cpu"
55
  model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
56
- token = os.environ['hf_token']
57
- pipe = pipeline("text-generation", model=model_name, device=device)
58
  generate_kwargs = {'max_new_tokens': 20}
59
 
60
  # '''
@@ -123,6 +122,7 @@ def generate(text, past_key_values):
123
  if __name__ == "__main__":
124
  with torch.no_grad():
125
  past_key_values = set_past_key_values()
 
126
  demo = gr.Interface(partial(generate, past_key_values=past_key_values),
127
  inputs="textbox", outputs="textbox")
128
  demo.launch()
 
53
  torch.set_grad_enabled(False)
54
  device = "cpu"
55
  model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
56
+ pipe = pipeline("text-generation", model=model_name, device='cuda')
 
57
  generate_kwargs = {'max_new_tokens': 20}
58
 
59
  # '''
 
122
  if __name__ == "__main__":
123
  with torch.no_grad():
124
  past_key_values = set_past_key_values()
125
+ pipe.model = pipe.model.cpu()
126
  demo = gr.Interface(partial(generate, past_key_values=past_key_values),
127
  inputs="textbox", outputs="textbox")
128
  demo.launch()