diff --git a/src/controller/mainFrameController.py b/src/controller/mainFrameController.py index 62d7fb6..cd06495 100644 --- a/src/controller/mainFrameController.py +++ b/src/controller/mainFrameController.py @@ -1,3 +1,4 @@ +from collections import deque import customtkinter as ctk from views.mainScreen import MainFrame from models.data import TextData @@ -5,6 +6,14 @@ from Ai.interence import VeraMindInference from utils.database.database import FakeNewsChecker from models.provider import Provider from collections import Counter +from Ai.llm import ArticleRater + +BAD_WORDS = ["FAKE", "SATIRE", "Fake", "fake"] +GOOD_WORDS = ["REAL", "real", "Real"] +BAD_COLOR = "#ff8080" +GOOD_COLOR = "#80ff8f" +WORDS = BAD_WORDS + GOOD_WORDS + class MainFrameController: """ @@ -22,6 +31,7 @@ class MainFrameController: self.model_inference = VeraMindInference('VeraMind-Mini') self.db = FakeNewsChecker() self.update_provider_list() + self.rater = ArticleRater() def get_text_data(self) -> TextData: """ @@ -35,16 +45,63 @@ class MainFrameController: text_data.text = self.frame.input_textbox.get("0.0", "end") return text_data - def press_check_button(self) -> None: - """ - Handle the 'Check' button press event. - Processes the input, makes a prediction, updates the database, and displays the result. - """ - text_data = self.get_text_data() - text_data = self._predict(text_data) - self._add_to_db(text_data) - self._update_output(text_data.get_output()) + def press_check_button(self): + text_data = self.get_textdata() + print(text_data.text) + self._predict(text_data) + self.frame.output_textbox.configure(state="normal") + self.frame.output_textbox.delete("0.0", "end") + + response_stream = self.rater.get_response(text_data.text, text_data.result, float(f"{text_data.confidence * 100:.2f}")) + + highlight_buffer = deque(maxlen=5) + + for chunk in response_stream: + # Display the chunk immediately + self.frame.output_textbox.insert("end", chunk) + self.frame.output_textbox.see("end") + self.frame.update_idletasks() + + # Add to highlight buffer + highlight_buffer.append(chunk) + + # Process highlighting when buffer is full + if len(highlight_buffer) == 5: + self.process_highlighting(highlight_buffer) + + # Process any remaining chunks in the buffer + if highlight_buffer: + self.process_highlighting(highlight_buffer) + + self.frame.output_textbox.configure(state="disabled") self.update_provider_list() + + def process_highlighting(self, highlight_buffer): + start_index = self.frame.output_textbox.index(f"end-{sum(len(c) for c in highlight_buffer)}c") + end_index = self.frame.output_textbox.index("end") + self.highlight_words(start_index, end_index) + + # Keep overlap of 2 chunks + highlight_buffer = deque(list(highlight_buffer)[-2:], maxlen=5) + + def highlight_words(self, start_index, end_index): + content = self.frame.output_textbox.get(start_index, end_index) + + for word in WORDS: + start = 0 + while True: + pos = content.find(word, start) + if pos == -1: + break + word_start = f"{start_index}+{pos}c" + word_end = f"{word_start}+{len(word)}c" + tag_name = f"{word.lower()}_color" + self.frame.output_textbox.tag_add(tag_name, word_start, word_end) + if word in BAD_WORDS: + self.frame.output_textbox.tag_config(tag_name, foreground=BAD_COLOR) + elif word in GOOD_WORDS: + self.frame.output_textbox.tag_config(tag_name, foreground=GOOD_COLOR) + start = pos + len(word) def _predict(self, text_data: TextData) -> TextData: """ @@ -57,8 +114,6 @@ class MainFrameController: text_data.confidence = result["confidence"] text_data.result = result["result"] text_data.is_fake_news = result["is_fake"] - print(f"Prediction: {text_data.result}") - print(f"Confidence: {text_data.confidence}") return text_data def _add_to_db(self, text_data: TextData) -> None: @@ -119,8 +174,6 @@ class MainFrameController: count_label = ctk.CTkLabel(provider_frame, text=str(provider.get_fake_percentage())+"%") count_label.pack(side="right", padx=5) - - def _update_output(self, output: str) -> None: """ Update the output text box with the result. @@ -130,6 +183,4 @@ class MainFrameController: self.frame.output_textbox.configure(state="normal") self.frame.output_textbox.delete("0.0", "end") self.frame.output_textbox.insert("0.0", output) - self.frame.output_textbox.configure(state="disabled") - - \ No newline at end of file + self.frame.output_textbox.configure(state="disabled") \ No newline at end of file