Spaces:
Sleeping
Sleeping
Update app.py
Browse files
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 |
-
|
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()
|