Spaces:
Runtime error
Runtime error
add support fur gpu
Browse files- app.R +1 -1
- model-session.R +12 -5
app.R
CHANGED
@@ -4,7 +4,7 @@ library(minhub)
|
|
4 |
library(magrittr)
|
5 |
source("model-session.R")
|
6 |
|
7 |
-
repo <- "
|
8 |
repo <- Sys.getenv("MODEL_REPO", unset = repo)
|
9 |
sess <- model_session$new()
|
10 |
|
|
|
4 |
library(magrittr)
|
5 |
source("model-session.R")
|
6 |
|
7 |
+
repo <- "stabilityai/stablelm-tuned-alpha-3b"
|
8 |
repo <- Sys.getenv("MODEL_REPO", unset = repo)
|
9 |
sess <- model_session$new()
|
10 |
|
model-session.R
CHANGED
@@ -18,9 +18,15 @@ model_session <- R6::R6Class(
|
|
18 |
library(torch)
|
19 |
library(zeallot)
|
20 |
library(minhub)
|
|
|
21 |
model <<- minhub::gptneox_from_pretrained(repo)
|
22 |
model$eval()
|
23 |
-
|
|
|
|
|
|
|
|
|
|
|
24 |
tok <<- tok::tokenizer$from_pretrained(repo)
|
25 |
"done"
|
26 |
})
|
@@ -36,15 +42,16 @@ model_session <- R6::R6Class(
|
|
36 |
top_k = self$top_k
|
37 |
)
|
38 |
self$task_q$push(args = args, function(prompt, temperature, top_k) {
|
39 |
-
|
|
|
40 |
with_no_grad({
|
41 |
-
logits <- model(idx + 1L)
|
42 |
})
|
43 |
logits <- logits[,-1,]/temperature
|
44 |
c(prob, ind) %<-% logits$topk(top_k)
|
45 |
-
logits <- torch_full_like(logits, -
|
46 |
logits <- nnf_softmax(logits, dim = -1)
|
47 |
-
id_next <- torch::torch_multinomial(logits, num_samples = 1) - 1L
|
48 |
tok$decode(as.integer(id_next))
|
49 |
})
|
50 |
}
|
|
|
18 |
library(torch)
|
19 |
library(zeallot)
|
20 |
library(minhub)
|
21 |
+
device <- if (cuda_is_available()) "cuda" else "cpu"
|
22 |
model <<- minhub::gptneox_from_pretrained(repo)
|
23 |
model$eval()
|
24 |
+
if (device == "cuda") {
|
25 |
+
model$to(device=device)
|
26 |
+
#model$to(dtype=torch_float())
|
27 |
+
} else {
|
28 |
+
model$to(dtype = torch_float())
|
29 |
+
}
|
30 |
tok <<- tok::tokenizer$from_pretrained(repo)
|
31 |
"done"
|
32 |
})
|
|
|
42 |
top_k = self$top_k
|
43 |
)
|
44 |
self$task_q$push(args = args, function(prompt, temperature, top_k) {
|
45 |
+
device <- if (cuda_is_available()) "cuda" else "cpu"
|
46 |
+
idx <- torch_tensor(tok$encode(prompt)$ids, device=device)$view(c(1, -1))
|
47 |
with_no_grad({
|
48 |
+
logits <- model(idx + 1L)$to(dtype="float", device="cpu")
|
49 |
})
|
50 |
logits <- logits[,-1,]/temperature
|
51 |
c(prob, ind) %<-% logits$topk(top_k)
|
52 |
+
logits <- torch_full_like(logits, -1e7)$scatter_(-1, ind, prob)
|
53 |
logits <- nnf_softmax(logits, dim = -1)
|
54 |
+
id_next <- torch::torch_multinomial(logits, num_samples = 1)$cpu() - 1L
|
55 |
tok$decode(as.integer(id_next))
|
56 |
})
|
57 |
}
|