Spaces:
Runtime error
Runtime error
use GPU
Browse files
app.py
CHANGED
@@ -16,10 +16,10 @@ example = [
|
|
16 |
[EXAMPLE_3, 11, 0.2, 42],
|
17 |
]
|
18 |
|
19 |
-
|
20 |
tokenizer = AutoTokenizer.from_pretrained("loubnabnl/santacoder-code-to-text")
|
21 |
-
model = AutoModelForCausalLM.from_pretrained("loubnabnl/santacoder-code-to-text", trust_remote_code=True)
|
22 |
-
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
|
23 |
|
24 |
def make_doctring(gen_prompt):
|
25 |
return gen_prompt + f"\n\n\"\"\"\nExplanation:"
|
|
|
16 |
[EXAMPLE_3, 11, 0.2, 42],
|
17 |
]
|
18 |
|
19 |
+
device="cuda:0"
|
20 |
tokenizer = AutoTokenizer.from_pretrained("loubnabnl/santacoder-code-to-text")
|
21 |
+
model = AutoModelForCausalLM.from_pretrained("loubnabnl/santacoder-code-to-text", trust_remote_code=True).to(device)
|
22 |
+
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
|
23 |
|
24 |
def make_doctring(gen_prompt):
|
25 |
return gen_prompt + f"\n\n\"\"\"\nExplanation:"
|