Sergidev commited on
Commit
9365e2e
1 Parent(s): 16dc17d

Update modules/pmbl.py

Browse files
Files changed (1) hide show
  1. modules/pmbl.py +6 -1
modules/pmbl.py CHANGED
@@ -1,8 +1,13 @@
1
  import sqlite3
2
  from datetime import datetime
3
- from llama_cpp import Llama
4
  from concurrent.futures import ThreadPoolExecutor
 
 
 
 
5
 
 
 
6
 
7
  class PMBL:
8
  def __init__(self, model_path):
 
1
  import sqlite3
2
  from datetime import datetime
 
3
  from concurrent.futures import ThreadPoolExecutor
4
+ import transformers
5
+
6
+ # CUDA support
7
+ transformers.install("llama-cpp-python", cuda=True)
8
 
9
+ #use llama_cpp
10
+ from llama_cpp import Llama
11
 
12
  class PMBL:
13
  def __init__(self, model_path):