65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
from celery import shared_task
|
|
from time import sleep
|
|
from redbeat import RedBeatSchedulerEntry
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.sql.expression import func, select, insert
|
|
from sqlalchemy.ext.automap import automap_base
|
|
from gpt4all import GPT4All
|
|
from sqlalchemy.orm import Session
|
|
from celery import current_app as celery_app
|
|
|
|
from .extensions import db
|
|
from .models import Result
|
|
|
|
engine = create_engine("sqlite:///words_prompts.db", pool_pre_ping=True)
|
|
|
|
Base = automap_base()
|
|
Base.prepare(engine, reflect=True)
|
|
|
|
Words = Base.classes.words
|
|
Themes = Base.classes.themes
|
|
|
|
MODEL = GPT4All(
|
|
model_name="gpt4all-falcon-q4_0.gguf",
|
|
# model_path=(Path.home() / ".cache" / "gpt4all"),
|
|
allow_download=False,
|
|
)
|
|
|
|
|
|
@shared_task
|
|
def my_task(text, schedule_name):
|
|
while True:
|
|
with Session(engine) as word_session:
|
|
random_word = word_session.query(Words.words)
|
|
random_word = random_word.order_by(func.random()).first()
|
|
random_word = str(random_word)[4:-4]
|
|
# SYSTEM_TEMPLATE = "A single sentence based on a word."
|
|
# PROMPT_TEMPLATE = "### Instruction: {0} \n### Response: "
|
|
response = MODEL.generate(
|
|
f"Give me a writing prompt about {random_word}.",
|
|
temp=0.7,
|
|
callback=stop_on_token_callback,
|
|
)
|
|
word_session.execute(insert(Themes).values(themes=response))
|
|
word_session.commit()
|
|
|
|
try:
|
|
entry = RedBeatSchedulerEntry.from_key(
|
|
"redbeat:" + schedule_name, app=celery_app
|
|
)
|
|
except KeyError:
|
|
entry = None
|
|
|
|
if entry:
|
|
entry.delete()
|
|
|
|
|
|
def stop_on_token_callback(token_id, token_string):
|
|
"""
|
|
Function to limit return length of the
|
|
gpt4all response. Period indicates a sentence.
|
|
"""
|
|
if "." in token_string:
|
|
return False
|
|
return True
|