NoteAI/text_processor.py at main · AlexWS12/NoteAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import re
from typing import List, Tuple
from config import CHUNK_SIZE_MIN, CHUNK_SIZE_MAX, CHUNK_OVERLAP_SENTENCES


class TextProcessor:

    @staticmethod
    def clean_text(text: str) -> str:
        text = re.sub(r' +', ' ', text)
        text = re.sub(r'\n\s*\n\s*\n+', '\n\n', text)
        text = re.sub(r'\n\s*\d+\s*\n', '\n', text)
        text = re.sub(r'Page \d+', '', text, flags=re.IGNORECASE)

        lines = text.split('\n')
        if len(lines) > 10:
            first_line = lines[0].strip()
            last_line = lines[-1].strip()
            if first_line and text.count(first_line) > 3:
                text = text.replace(first_line, '', 1)
            if last_line and text.count(last_line) > 3:
                text = text[:text.rfind(last_line)]

        return text.strip()

    @staticmethod
    def extract_text_from_file(file_content: bytes, file_type: str) -> str:
        if file_type == "txt":
            return file_content.decode('utf-8', errors='ignore')
        elif file_type == "pdf":
            try:
                from PyPDF2 import PdfReader
                from io import BytesIO

                pdf_file = BytesIO(file_content)
                pdf_reader = PdfReader(pdf_file)

                text = ""
                for page in pdf_reader.pages:
                    text += page.extract_text() + "\n"

                return text
            except Exception as e:
                raise Exception(f"PDF extraction failed: {str(e)}")
        else:
            raise ValueError(f"Unsupported file type: {file_type}")

    @staticmethod
    def chunk_text(text: str, max_chunk_size: int = CHUNK_SIZE_MAX) -> List[str]:
        words = text.split()

        if len(words) <= max_chunk_size:
            return [text]

        chunks = []
        sentences = re.split(r'[.!?]+', text)
        sentences = [s.strip() for s in sentences if s.strip()]

        current_chunk = []
        current_word_count = 0

        for sentence in sentences:
            sentence_words = sentence.split()
            sentence_word_count = len(sentence_words)

            if current_word_count + sentence_word_count > max_chunk_size and current_chunk:
                chunks.append(' '.join(current_chunk))

                overlap_sentences = current_chunk[-CHUNK_OVERLAP_SENTENCES:] if len(current_chunk) >= CHUNK_OVERLAP_SENTENCES else current_chunk
                current_chunk = overlap_sentences
                current_word_count = sum(len(s.split()) for s in overlap_sentences)

            current_chunk.append(sentence)
            current_word_count += sentence_word_count

        if current_chunk:
            chunks.append(' '.join(current_chunk))

        return chunks

    @staticmethod
    def count_words(text: str) -> int:
        return len(text.split())

    @staticmethod
    def validate_input(text: str, max_length: int) -> Tuple[bool, str]:
        if not text or not text.strip():
            return False, "Text cannot be empty"

        if len(text) > max_length:
            return False, f"Text exceeds maximum length of {max_length} characters"

        return True, ""