transformer: read input from unix socket ()
This commit is contained in:
parent
5aba6d97ec
commit
480cec101f
47
transformer.py
Normal file
47
transformer.py
Normal 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()
|
Loading…
Reference in New Issue
Block a user