undo cuda
Browse files- modules/pmbl.py +3 -8
modules/pmbl.py
CHANGED
@@ -1,7 +1,6 @@
|
|
1 |
import sqlite3
|
2 |
from datetime import datetime
|
3 |
from concurrent.futures import ThreadPoolExecutor
|
4 |
-
from llama_cpp import Llama
|
5 |
|
6 |
class PMBL:
|
7 |
def __init__(self, model_path):
|
@@ -101,11 +100,9 @@ class PMBL:
|
|
101 |
|
102 |
for chunk in response.result():
|
103 |
yield chunk
|
104 |
-
|
105 |
-
|
106 |
def generate_response_task(self, system_prompt, prompt, n_ctx):
|
107 |
llm = Llama(model_path=self.model_path, n_ctx=n_ctx, n_threads=8, mlock=True)
|
108 |
-
llm = llm.to("cuda") # Move the model to the GPU
|
109 |
|
110 |
response = llm(
|
111 |
system_prompt,
|
@@ -149,11 +146,9 @@ class PMBL:
|
|
149 |
conn.commit()
|
150 |
|
151 |
conn.close()
|
152 |
-
|
153 |
-
|
154 |
def generate_topic(self, prompt, response):
|
155 |
llm = Llama(model_path=self.model_path, n_ctx=1690, n_threads=8, mlock=True)
|
156 |
-
llm = llm.to("cuda") # Move the model to the GPU
|
157 |
|
158 |
system_prompt = f"Based on the following interaction between a user and an AI assistant, generate a concise topic for the conversation in 2-4 words:\n\nUser: {prompt}\nAssistant: {response}\n\nTopic:"
|
159 |
|
@@ -165,4 +160,4 @@ class PMBL:
|
|
165 |
echo=False
|
166 |
)
|
167 |
|
168 |
-
return topic['choices'][0]['text'].strip()
|
|
|
1 |
import sqlite3
|
2 |
from datetime import datetime
|
3 |
from concurrent.futures import ThreadPoolExecutor
|
|
|
4 |
|
5 |
class PMBL:
|
6 |
def __init__(self, model_path):
|
|
|
100 |
|
101 |
for chunk in response.result():
|
102 |
yield chunk
|
103 |
+
|
|
|
104 |
def generate_response_task(self, system_prompt, prompt, n_ctx):
|
105 |
llm = Llama(model_path=self.model_path, n_ctx=n_ctx, n_threads=8, mlock=True)
|
|
|
106 |
|
107 |
response = llm(
|
108 |
system_prompt,
|
|
|
146 |
conn.commit()
|
147 |
|
148 |
conn.close()
|
149 |
+
|
|
|
150 |
def generate_topic(self, prompt, response):
|
151 |
llm = Llama(model_path=self.model_path, n_ctx=1690, n_threads=8, mlock=True)
|
|
|
152 |
|
153 |
system_prompt = f"Based on the following interaction between a user and an AI assistant, generate a concise topic for the conversation in 2-4 words:\n\nUser: {prompt}\nAssistant: {response}\n\nTopic:"
|
154 |
|
|
|
160 |
echo=False
|
161 |
)
|
162 |
|
163 |
+
return topic['choices'][0]['text'].strip()
|