Files
crowsnest/python/ai_api/main.py

52 lines
1.4 KiB
Python

from transformers import pipeline
import psycopg2
# Connect to the School database
def summarize_articles(summfunc):
""" Connect to the PostgreSQL database server """
conn = psycopg2.connect(
dbname="crowsnest",
user="crow",
password="4LlKpnQ2RZPzL13BSpkW4k",
host="192.168.0.2"
)
try:
# create a cursor
cur = conn.cursor()
for i in range(20):
print(i)
# execute a statement
cur.execute("""
SELECT id, content
FROM articles
WHERE aisummary = ''
ORDER BY publishDate DESC
LIMIT 20
""")
# display the PostgreSQL database server version
keys, contents = list(zip(*cur.fetchall()))
for key, summary in zip(keys, summfunc(list(contents))):
cur.execute("UPDATE articles SET aisummary = %s WHERE id = %s", (summary, key))
conn.commit()
# close the communication with the PostgreSQL
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
raise error
finally:
if conn is not None:
conn.close()
print('Database connection closed.')
if __name__ == '__main__':
pipe = pipeline("summarization", model="Shahm/bart-german")
summarize_articles(lambda x: [str(e["summary_text"]) for e in pipe(x)])