SemanticTest

1
from sentence_transformers import SentenceTransformer
2
import socket
3
import json
4
import select
5
import sys
6
import win32com.client
7
import pypdfium2 as pdfium
8
import textwrap as tw
9

10
exts_docs = ['.docx', '.doc', '.docm', '.dotx', '.dot', '.dotm', '.rtf']
11
exts_pdf = ['.pdf']
12
exts_txt = ['.txt']
13

14
encodings = ['utf-8', 'ansi', 'utf-8-sig', 'utf-16', 'utf-32', 'cp1252', 'cp1250', 'cp1251', 'latin-1', 'ascii', 'iso-8859-1', 'utf-32']
15

16
def main():
17
    if len(sys.argv) < 2:
18
        return
19
    print(sys.argv[1])
20

21
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
22
    server_socket.bind(('127.0.0.1', 19200))
23
    server_socket.listen(5)
24
    print(f"Сервер запущен на порту 19200.")
25
    client_socket, client_address = server_socket.accept()
26

27
    try:
28
        model_name = 'sentence-transformers/LaBSE'
29
        model = SentenceTransformer(model_name, cache_folder=f"{sys.argv[1]}")
30
        load = model.encode("loading script")
31
        client_socket.send(b'Loaded\n')
32
    except:
33
        client_socket.send(b'Error\n')
34
        exit()
35
    print(f"Подключение с {client_address} установлено.")
36

37
    print(f'Загружена модель {model_name}.')
38

39
    def ends_for(string, exts):
40
        for ext in exts:
41
            if string.endswith(ext):
42
                return True
43
        return False
44

45
    def check_for_encoding(filename, encoding):
46
        try:
47
            with open(filename, 'r', encoding=encoding) as f:
48
                text = f.read()
49
                return True, text
50
        except UnicodeDecodeError:
51
            return False, ""
52

53
    def index_file(filename):
54

55
        if ends_for(filename, exts_txt):
56
            for encoding in encodings:
57
                checked, text = check_for_encoding(filename, encoding)
58
                if not checked:
59
                    continue
60

61
                text_embeddings = []
62
                tw_text = tw.wrap(text, 10000, break_long_words=False)
63
                for piece in tw_text:
64
                    text_embeddings.append(model.encode(piece).tolist())
65
                return text_embeddings
66

67
        elif ends_for(filename, exts_docs):
68
            doc = win32com.client.GetObject(filename)
69
            text = doc.Content.Text
70
            text1 = doc.Range().Text
71
            print(text == text1)
72
            text_embeddings = []
73
            tw_text = tw.wrap(text, 10000, break_long_words=False)
74
            for piece in tw_text:
75
                text_embeddings.append(model.encode(piece).tolist())
76
            return text_embeddings
77

78
        elif ends_for(filename, exts_pdf):
79
            pdf = pdfium.PdfDocument(filename)
80
            text = []
81
            text_embeddings = []
82
            for page in pdf:
83
                text.append(page.get_textpage().get_text_range() + '\n')
84
            tw_text = tw.wrap('\n'.join(text), 10000, break_long_words=False)
85
            for piece in tw_text:
86
                text_embeddings.append(model.encode(piece).tolist())
87
            return text_embeddings
88

89
    def process_query():
90
        client_socket.setblocking(False)
91
        while 1 > 0:
92
            ready = select.select([client_socket], [], [], 1)
93
            if ready[0]:
94
                try:
95
                    data = socket.SocketIO(sock=client_socket, mode='r').readline()
96
                    if not data:
97
                        break
98

99
                    json_data = json.loads(data.decode('utf-8'))
100

101
                    if json_data['Type'] == 'Close':
102
                        print(f"Закрытие соединения с {client_address}.")
103
                        client_socket.close()
104
                        server_socket.close()
105
                        break
106
                    elif json_data['Type'] == 'IndexingFile':
107
                        print(f"Получен запрос на индексирование файла {json_data['Type']}")
108
                        embeddings = index_file(str(json_data['Value']))
109
                        client_socket.send(f"{json.dumps(embeddings)}\n".encode('utf-8'))
110
                    elif json_data['Type'] == 'IndexingText':
111
                        print(f"Получен запрос на индексирование текста {json_data['Type']}")
112
                        embeddings = model.encode(str(json_data['Value']))
113
                        client_socket.send(f"{json.dumps(embeddings.tolist())}\n".encode('utf-8'))
114
                except:
115
                    print(f"Ошибка при получении данных")
116
                    break
117

118
    process_query()
119

120

121
main()
122