transformer: read input from unix socket ()

This commit is contained in:
Felix Pankratz 2022-07-21 21:14:02 +02:00
parent 5aba6d97ec
commit 480cec101f

47
transformer.py Normal file
View File

@ -0,0 +1,47 @@
#!/usr/bin/env python3
from transformers import pipeline
import sys
import socket
import os
def main():
#socket_addr = '/mnt/d/devel/hotpulp/hp'
socket_addr = '/home/panki/hp'
try:
os.unlink(socket_addr)
except:
if os.path.exists(socket_addr):
raise
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
s.bind(socket_addr)
s.listen(1)
print('Socket ready: Listening on {}'.format(socket_addr))
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
print('Summarizer ready')
while True:
connection, client_address = s.accept()
try:
buffer = ''
rcvd_bytes = 0
while True:
data = connection.recv(1024)
rcvd_bytes = rcvd_bytes + len(data)
if not data:
print('Finished receiving data, processing...')
break
buffer = buffer + str(data)
buffer = buffer[:-1].rstrip('\x00')
try:
print(summarizer(buffer, max_length=500, min_length=100, do_sample=False))
except IndexError:
print("ERROR: Input too large!")
finally:
connection.close()
print('Connection closed')
if __name__ == '__main__':
main()