mirror of
https://github.com/Monadical-SAS/reflector.git
synced 2025-12-20 12:19:06 +00:00
265 lines
8.6 KiB
Python
265 lines
8.6 KiB
Python
"""
|
|
Utility file for all text processing related functionalities
|
|
"""
|
|
import datetime
|
|
from typing import List
|
|
|
|
import nltk
|
|
import torch
|
|
from nltk.corpus import stopwords
|
|
from nltk.tokenize import word_tokenize
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
from transformers import BartForConditionalGeneration, BartTokenizer
|
|
|
|
from log_utils import LOGGER
|
|
from run_utils import CONFIG
|
|
|
|
nltk.download("punkt", quiet=True)
|
|
|
|
|
|
def preprocess_sentence(sentence: str) -> str:
|
|
"""
|
|
Filter out undesirable tokens from thr sentence
|
|
:param sentence:
|
|
:return:
|
|
"""
|
|
stop_words = set(stopwords.words("english"))
|
|
tokens = word_tokenize(sentence.lower())
|
|
tokens = [token for token in tokens if token.isalnum() and token not in stop_words]
|
|
return " ".join(tokens)
|
|
|
|
|
|
def compute_similarity(sent1: str, sent2: str) -> float:
|
|
"""
|
|
Compute the similarity
|
|
"""
|
|
tfidf_vectorizer = TfidfVectorizer()
|
|
if sent1 is not None and sent2 is not None:
|
|
tfidf_matrix = tfidf_vectorizer.fit_transform([sent1, sent2])
|
|
return cosine_similarity(tfidf_matrix[0], tfidf_matrix[1])[0][0]
|
|
return 0.0
|
|
|
|
|
|
def remove_almost_alike_sentences(sentences: List[str], threshold=0.7) -> List[str]:
|
|
"""
|
|
Filter sentences that are similar beyond a set threshold
|
|
:param sentences:
|
|
:param threshold:
|
|
:return:
|
|
"""
|
|
num_sentences = len(sentences)
|
|
removed_indices = set()
|
|
|
|
for i in range(num_sentences):
|
|
if i not in removed_indices:
|
|
for j in range(i + 1, num_sentences):
|
|
if j not in removed_indices:
|
|
l_i = len(sentences[i])
|
|
l_j = len(sentences[j])
|
|
if l_i == 0 or l_j == 0:
|
|
if l_i == 0:
|
|
removed_indices.add(i)
|
|
if l_j == 0:
|
|
removed_indices.add(j)
|
|
else:
|
|
sentence1 = preprocess_sentence(sentences[i])
|
|
sentence2 = preprocess_sentence(sentences[j])
|
|
if len(sentence1) != 0 and len(sentence2) != 0:
|
|
similarity = compute_similarity(sentence1, sentence2)
|
|
|
|
if similarity >= threshold:
|
|
removed_indices.add(max(i, j))
|
|
|
|
filtered_sentences = [
|
|
sentences[i] for i in range(num_sentences) if i not in removed_indices
|
|
]
|
|
return filtered_sentences
|
|
|
|
|
|
def remove_outright_duplicate_sentences_from_chunk(chunk: str) -> List[str]:
|
|
"""
|
|
Remove repetitive sentences
|
|
:param chunk:
|
|
:return:
|
|
"""
|
|
chunk_text = chunk["text"]
|
|
sentences = nltk.sent_tokenize(chunk_text)
|
|
nonduplicate_sentences = list(dict.fromkeys(sentences))
|
|
return nonduplicate_sentences
|
|
|
|
|
|
def remove_whisper_repetitive_hallucination(
|
|
nonduplicate_sentences: List[str],
|
|
) -> List[str]:
|
|
"""
|
|
Remove sentences that are repeated as a result of Whisper
|
|
hallucinations
|
|
:param nonduplicate_sentences:
|
|
:return:
|
|
"""
|
|
chunk_sentences = []
|
|
|
|
for sent in nonduplicate_sentences:
|
|
temp_result = ""
|
|
seen = {}
|
|
words = nltk.word_tokenize(sent)
|
|
n_gram_filter = 3
|
|
for i in range(len(words)):
|
|
if (
|
|
str(words[i : i + n_gram_filter]) in seen
|
|
and seen[str(words[i : i + n_gram_filter])]
|
|
== words[i + 1 : i + n_gram_filter + 2]
|
|
):
|
|
pass
|
|
else:
|
|
seen[str(words[i : i + n_gram_filter])] = words[
|
|
i + 1 : i + n_gram_filter + 2
|
|
]
|
|
temp_result += words[i]
|
|
temp_result += " "
|
|
chunk_sentences.append(temp_result)
|
|
return chunk_sentences
|
|
|
|
|
|
def post_process_transcription(whisper_result: dict) -> dict:
|
|
"""
|
|
Parent function to perform post-processing on the transcription result
|
|
:param whisper_result:
|
|
:return:
|
|
"""
|
|
transcript_text = ""
|
|
for chunk in whisper_result["chunks"]:
|
|
nonduplicate_sentences = remove_outright_duplicate_sentences_from_chunk(chunk)
|
|
chunk_sentences = remove_whisper_repetitive_hallucination(
|
|
nonduplicate_sentences
|
|
)
|
|
similarity_matched_sentences = remove_almost_alike_sentences(chunk_sentences)
|
|
chunk["text"] = " ".join(similarity_matched_sentences)
|
|
transcript_text += chunk["text"]
|
|
whisper_result["text"] = transcript_text
|
|
return whisper_result
|
|
|
|
|
|
def summarize_chunks(chunks: List[str], tokenizer, model) -> List[str]:
|
|
"""
|
|
Summarize each chunk using a summarizer model
|
|
:param chunks:
|
|
:param tokenizer:
|
|
:param model:
|
|
:return:
|
|
"""
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
summaries = []
|
|
for c in chunks:
|
|
input_ids = tokenizer.encode(c, return_tensors="pt")
|
|
input_ids = input_ids.to(device)
|
|
with torch.no_grad():
|
|
summary_ids = model.generate(
|
|
input_ids,
|
|
num_beams=int(CONFIG["SUMMARIZER"]["BEAM_SIZE"]),
|
|
length_penalty=2.0,
|
|
max_length=int(CONFIG["SUMMARIZER"]["MAX_LENGTH"]),
|
|
early_stopping=True,
|
|
)
|
|
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
|
|
summaries.append(summary)
|
|
return summaries
|
|
|
|
|
|
def chunk_text(
|
|
text: str, max_chunk_length: int = int(CONFIG["SUMMARIZER"]["MAX_CHUNK_LENGTH"])
|
|
) -> List[str]:
|
|
"""
|
|
Split text into smaller chunks.
|
|
:param text: Text to be chunked
|
|
:param max_chunk_length: length of chunk
|
|
:return: chunked texts
|
|
"""
|
|
sentences = nltk.sent_tokenize(text)
|
|
chunks = []
|
|
current_chunk = ""
|
|
for sentence in sentences:
|
|
if len(current_chunk) + len(sentence) < max_chunk_length:
|
|
current_chunk += f" {sentence.strip()}"
|
|
else:
|
|
chunks.append(current_chunk.strip())
|
|
current_chunk = f"{sentence.strip()}"
|
|
chunks.append(current_chunk.strip())
|
|
return chunks
|
|
|
|
|
|
def summarize(
|
|
transcript_text: str,
|
|
timestamp: datetime.datetime.timestamp,
|
|
real_time: bool = False,
|
|
chunk_summarize: str = CONFIG["SUMMARIZER"]["SUMMARIZE_USING_CHUNKS"],
|
|
):
|
|
"""
|
|
Summarize the given text either as a whole or as chunks as needed
|
|
:param transcript_text:
|
|
:param timestamp:
|
|
:param real_time:
|
|
:param chunk_summarize:
|
|
:return:
|
|
"""
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
summary_model = CONFIG["SUMMARIZER"]["SUMMARY_MODEL"]
|
|
if not summary_model:
|
|
summary_model = "facebook/bart-large-cnn"
|
|
|
|
# Summarize the generated transcript using the BART model
|
|
LOGGER.info(f"Loading BART model: {summary_model}")
|
|
tokenizer = BartTokenizer.from_pretrained(summary_model)
|
|
model = BartForConditionalGeneration.from_pretrained(summary_model)
|
|
model = model.to(device)
|
|
|
|
output_file = "summary_" + timestamp.strftime("%m-%d-%Y_%H:%M:%S") + ".txt"
|
|
if real_time:
|
|
output_file = "real_time_" + output_file
|
|
|
|
if chunk_summarize != "YES":
|
|
max_length = int(CONFIG["SUMMARIZER"]["INPUT_ENCODING_MAX_LENGTH"])
|
|
inputs = tokenizer.batch_encode_plus(
|
|
[transcript_text],
|
|
truncation=True,
|
|
padding="longest",
|
|
max_length=max_length,
|
|
return_tensors="pt",
|
|
)
|
|
inputs = inputs.to(device)
|
|
|
|
with torch.no_grad():
|
|
num_beans = int(CONFIG["SUMMARIZER"]["BEAM_SIZE"])
|
|
max_length = int(CONFIG["SUMMARIZER"]["MAX_LENGTH"])
|
|
summaries = model.generate(
|
|
inputs["input_ids"],
|
|
num_beams=num_beans,
|
|
length_penalty=2.0,
|
|
max_length=max_length,
|
|
early_stopping=True,
|
|
)
|
|
|
|
decoded_summaries = [
|
|
tokenizer.decode(
|
|
summary, skip_special_tokens=True, clean_up_tokenization_spaces=False
|
|
)
|
|
for summary in summaries
|
|
]
|
|
summary = " ".join(decoded_summaries)
|
|
with open("./artefacts/" + output_file, "w", encoding="utf-8") as file:
|
|
file.write(summary.strip() + "\n")
|
|
else:
|
|
LOGGER.info("Breaking transcript into smaller chunks")
|
|
chunks = chunk_text(transcript_text)
|
|
|
|
LOGGER.info(
|
|
f"Transcript broken into {len(chunks)} " f"chunks of at most 500 words"
|
|
)
|
|
|
|
LOGGER.info(f"Writing summary text to: {output_file}")
|
|
with open(output_file, "w") as f:
|
|
summaries = summarize_chunks(chunks, tokenizer, model)
|
|
for summary in summaries:
|
|
f.write(summary.strip() + " ")
|