oceansweep commited on
Commit
43cd37c
1 Parent(s): cd34bd3

Upload 155 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. App_Function_Libraries/Audio/Audio_Files.py +786 -0
  2. App_Function_Libraries/Audio/Audio_Transcription_Lib.py +335 -0
  3. App_Function_Libraries/Audio/Diarization_Lib.py +275 -0
  4. App_Function_Libraries/Audio/__init__.py +0 -0
  5. App_Function_Libraries/Benchmarks_Evaluations/Confabulation_check.py +81 -0
  6. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/.gitignore +5 -0
  7. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/LICENSE +23 -0
  8. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/__init__.py +0 -0
  9. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/config.txt +30 -0
  10. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_multi_api.py +300 -0
  11. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_utils.py +730 -0
  12. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/prompt.py +62 -0
  13. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/test_chat_API_Calls.py +106 -0
  14. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README.md +200 -0
  15. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README_ZH.md +172 -0
  16. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/__init__.py +0 -0
  17. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/InfiniteBench/PUT_DATASETS_HERE.txt +0 -0
  18. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/__init__.py +0 -0
  19. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/collections.json +1 -0
  20. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/construct_synthetic_dataset.py +413 -0
  21. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/functions_module.py +1650 -0
  22. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/requirements.txt +9 -0
  23. App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/scripts/download_dataset.sh +6 -0
  24. App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/MMLU_Pro_rewritten.py +341 -0
  25. App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/__init__.py +0 -0
  26. App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/config.toml +30 -0
  27. App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/mmlu_pro_test.py +232 -0
  28. App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/run_openai.py +546 -0
  29. App_Function_Libraries/Benchmarks_Evaluations/__init__.py +0 -0
  30. App_Function_Libraries/Benchmarks_Evaluations/ms_g_eval.py +498 -0
  31. App_Function_Libraries/Books/.pytest_cache/.gitignore +2 -0
  32. App_Function_Libraries/Books/.pytest_cache/CACHEDIR.TAG +4 -0
  33. App_Function_Libraries/Books/.pytest_cache/README.md +8 -0
  34. App_Function_Libraries/Books/.pytest_cache/v/cache/lastfailed +10 -0
  35. App_Function_Libraries/Books/.pytest_cache/v/cache/nodeids +11 -0
  36. App_Function_Libraries/Books/.pytest_cache/v/cache/stepwise +1 -0
  37. App_Function_Libraries/Books/Book_Ingestion_Lib.py +577 -0
  38. App_Function_Libraries/Books/__init__.py +0 -0
  39. App_Function_Libraries/Character_Chat/Character_Chat_Lib.py +607 -0
  40. App_Function_Libraries/Character_Chat/__init__.py +0 -0
  41. App_Function_Libraries/Chat.py +439 -0
  42. App_Function_Libraries/Chunk_Lib.py +1051 -0
  43. App_Function_Libraries/DB/Character_Chat_DB.py +701 -0
  44. App_Function_Libraries/DB/DB_Manager.py +991 -0
  45. App_Function_Libraries/DB/RAG_QA_Chat_DB.py +722 -0
  46. App_Function_Libraries/DB/SQLite_DB.py +0 -0
  47. App_Function_Libraries/DB/__init__.py +0 -0
  48. App_Function_Libraries/Gradio_Related.py +420 -0
  49. App_Function_Libraries/Gradio_UI/Arxiv_tab.py +230 -0
  50. App_Function_Libraries/Gradio_UI/Audio_ingestion_tab.py +167 -0
App_Function_Libraries/Audio/Audio_Files.py ADDED
@@ -0,0 +1,786 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Audio_Files.py
2
+ #########################################
3
+ # Audio Processing Library
4
+ # This library is used to download or load audio files from a local directory.
5
+ #
6
+ ####
7
+ #
8
+ # Functions:
9
+ #
10
+ # download_audio_file(url, save_path)
11
+ # process_audio(
12
+ # process_audio_file(audio_url, audio_file, whisper_model="small.en", api_name=None, api_key=None)
13
+ #
14
+ #
15
+ #########################################
16
+ # Imports
17
+ import json
18
+ import logging
19
+ import os
20
+ import subprocess
21
+ import tempfile
22
+ import time
23
+ import uuid
24
+ from datetime import datetime
25
+ from pathlib import Path
26
+ #
27
+ # External Imports
28
+ import requests
29
+ import yt_dlp
30
+ #
31
+ # Local Imports
32
+ from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords, \
33
+ check_media_and_whisper_model
34
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
35
+ from App_Function_Libraries.Summarization.Summarization_General_Lib import perform_summarization
36
+ from App_Function_Libraries.Utils.Utils import downloaded_files, \
37
+ sanitize_filename, generate_unique_id, temp_files
38
+ from App_Function_Libraries.Video_DL_Ingestion_Lib import extract_metadata
39
+ from App_Function_Libraries.Audio.Audio_Transcription_Lib import speech_to_text
40
+ from App_Function_Libraries.Chunk_Lib import improved_chunking_process
41
+ #
42
+ #######################################################################################################################
43
+ # Function Definitions
44
+ #
45
+
46
+ MAX_FILE_SIZE = 500 * 1024 * 1024
47
+
48
+
49
+ def download_audio_file(url, current_whisper_model="", use_cookies=False, cookies=None):
50
+ try:
51
+ # Check if media already exists in the database and compare whisper models
52
+ should_download, reason = check_media_and_whisper_model(
53
+ url=url,
54
+ current_whisper_model=current_whisper_model
55
+ )
56
+
57
+ if not should_download:
58
+ logging.info(f"Skipping audio download: {reason}")
59
+ return None
60
+
61
+ logging.info(f"Proceeding with audio download: {reason}")
62
+
63
+ # Set up the request headers
64
+ headers = {}
65
+ if use_cookies and cookies:
66
+ try:
67
+ cookie_dict = json.loads(cookies)
68
+ headers['Cookie'] = '; '.join([f'{k}={v}' for k, v in cookie_dict.items()])
69
+ except json.JSONDecodeError:
70
+ logging.warning("Invalid cookie format. Proceeding without cookies.")
71
+
72
+ # Make the request
73
+ response = requests.get(url, headers=headers, stream=True)
74
+ # Raise an exception for bad status codes
75
+ response.raise_for_status()
76
+
77
+ # Get the file size
78
+ file_size = int(response.headers.get('content-length', 0))
79
+ if file_size > 500 * 1024 * 1024: # 500 MB limit
80
+ raise ValueError("File size exceeds the 500MB limit.")
81
+
82
+ # Generate a unique filename
83
+ file_name = f"audio_{uuid.uuid4().hex[:8]}.mp3"
84
+ save_path = os.path.join('downloads', file_name)
85
+
86
+ # Ensure the downloads directory exists
87
+ os.makedirs('downloads', exist_ok=True)
88
+
89
+
90
+ # Download the file
91
+ with open(save_path, 'wb') as f:
92
+ for chunk in response.iter_content(chunk_size=8192):
93
+ if chunk:
94
+ f.write(chunk)
95
+
96
+ logging.info(f"Audio file downloaded successfully: {save_path}")
97
+ return save_path
98
+
99
+ except requests.RequestException as e:
100
+ logging.error(f"Error downloading audio file: {str(e)}")
101
+ raise
102
+ except ValueError as e:
103
+ logging.error(str(e))
104
+ raise
105
+ except Exception as e:
106
+ logging.error(f"Unexpected error downloading audio file: {str(e)}")
107
+ raise
108
+
109
+ def process_audio_files(audio_urls, audio_file, whisper_model, api_name, api_key, use_cookies, cookies, keep_original,
110
+ custom_keywords, custom_prompt_input, chunk_method, max_chunk_size, chunk_overlap,
111
+ use_adaptive_chunking, use_multi_level_chunking, chunk_language, diarize,
112
+ keep_timestamps, custom_title):
113
+
114
+ start_time = time.time() # Start time for processing
115
+ processed_count = 0
116
+ failed_count = 0
117
+ progress = []
118
+ all_transcriptions = []
119
+ all_summaries = []
120
+ #v2
121
+ def format_transcription_with_timestamps(segments):
122
+ if keep_timestamps:
123
+ formatted_segments = []
124
+ for segment in segments:
125
+ start = segment.get('Time_Start', 0)
126
+ end = segment.get('Time_End', 0)
127
+ text = segment.get('Text', '').strip() # Ensure text is stripped of leading/trailing spaces
128
+
129
+ # Add the formatted timestamp and text to the list, followed by a newline
130
+ formatted_segments.append(f"[{start:.2f}-{end:.2f}] {text}")
131
+
132
+ # Join the segments with a newline to ensure proper formatting
133
+ return "\n".join(formatted_segments)
134
+ else:
135
+ # Join the text without timestamps
136
+ return "\n".join([segment.get('Text', '').strip() for segment in segments])
137
+
138
+ def update_progress(message):
139
+ progress.append(message)
140
+ return "\n".join(progress)
141
+
142
+ def cleanup_files():
143
+ for file in temp_files:
144
+ try:
145
+ if os.path.exists(file):
146
+ os.remove(file)
147
+ update_progress(f"Temporary file {file} removed.")
148
+ except Exception as e:
149
+ update_progress(f"Failed to remove temporary file {file}: {str(e)}")
150
+
151
+ def reencode_mp3(mp3_file_path):
152
+ try:
153
+ reencoded_mp3_path = mp3_file_path.replace(".mp3", "_reencoded.mp3")
154
+ subprocess.run([ffmpeg_cmd, '-i', mp3_file_path, '-codec:a', 'libmp3lame', reencoded_mp3_path], check=True)
155
+ update_progress(f"Re-encoded {mp3_file_path} to {reencoded_mp3_path}.")
156
+ return reencoded_mp3_path
157
+ except subprocess.CalledProcessError as e:
158
+ update_progress(f"Error re-encoding {mp3_file_path}: {str(e)}")
159
+ raise
160
+
161
+ def convert_mp3_to_wav(mp3_file_path):
162
+ try:
163
+ wav_file_path = mp3_file_path.replace(".mp3", ".wav")
164
+ subprocess.run([ffmpeg_cmd, '-i', mp3_file_path, wav_file_path], check=True)
165
+ update_progress(f"Converted {mp3_file_path} to {wav_file_path}.")
166
+ return wav_file_path
167
+ except subprocess.CalledProcessError as e:
168
+ update_progress(f"Error converting {mp3_file_path} to WAV: {str(e)}")
169
+ raise
170
+
171
+ try:
172
+ # Check and set the ffmpeg command
173
+ global ffmpeg_cmd
174
+ if os.name == "nt":
175
+ logging.debug("Running on Windows")
176
+ ffmpeg_cmd = os.path.join(os.getcwd(), "Bin", "ffmpeg.exe")
177
+ else:
178
+ ffmpeg_cmd = 'ffmpeg' # Assume 'ffmpeg' is in PATH for non-Windows systems
179
+
180
+ # Ensure ffmpeg is accessible
181
+ if not os.path.exists(ffmpeg_cmd) and os.name == "nt":
182
+ raise FileNotFoundError(f"ffmpeg executable not found at path: {ffmpeg_cmd}")
183
+
184
+ # Define chunk options early to avoid undefined errors
185
+ chunk_options = {
186
+ 'method': chunk_method,
187
+ 'max_size': max_chunk_size,
188
+ 'overlap': chunk_overlap,
189
+ 'adaptive': use_adaptive_chunking,
190
+ 'multi_level': use_multi_level_chunking,
191
+ 'language': chunk_language
192
+ }
193
+
194
+ # Process multiple URLs
195
+ urls = [url.strip() for url in audio_urls.split('\n') if url.strip()]
196
+
197
+ for i, url in enumerate(urls):
198
+ update_progress(f"Processing URL {i + 1}/{len(urls)}: {url}")
199
+
200
+ # Download and process audio file
201
+ audio_file_path = download_audio_file(url, use_cookies, cookies)
202
+ if not os.path.exists(audio_file_path):
203
+ update_progress(f"Downloaded file not found: {audio_file_path}")
204
+ failed_count += 1
205
+ log_counter(
206
+ metric_name="audio_files_failed_total",
207
+ labels={"whisper_model": whisper_model, "api_name": api_name},
208
+ value=1
209
+ )
210
+ continue
211
+
212
+ temp_files.append(audio_file_path)
213
+ update_progress("Audio file downloaded successfully.")
214
+
215
+ # Re-encode MP3 to fix potential issues
216
+ reencoded_mp3_path = reencode_mp3(audio_file_path)
217
+ if not os.path.exists(reencoded_mp3_path):
218
+ update_progress(f"Re-encoded file not found: {reencoded_mp3_path}")
219
+ failed_count += 1
220
+ log_counter(
221
+ metric_name="audio_files_failed_total",
222
+ labels={"whisper_model": whisper_model, "api_name": api_name},
223
+ value=1
224
+ )
225
+ continue
226
+
227
+ temp_files.append(reencoded_mp3_path)
228
+
229
+ # Convert re-encoded MP3 to WAV
230
+ wav_file_path = convert_mp3_to_wav(reencoded_mp3_path)
231
+ if not os.path.exists(wav_file_path):
232
+ update_progress(f"Converted WAV file not found: {wav_file_path}")
233
+ failed_count += 1
234
+ log_counter(
235
+ metric_name="audio_files_failed_total",
236
+ labels={"whisper_model": whisper_model, "api_name": api_name},
237
+ value=1
238
+ )
239
+ continue
240
+
241
+ temp_files.append(wav_file_path)
242
+
243
+ # Initialize transcription
244
+ transcription = ""
245
+
246
+ # Transcribe audio
247
+ if diarize:
248
+ segments = speech_to_text(wav_file_path, whisper_model=whisper_model, diarize=True)
249
+ else:
250
+ segments = speech_to_text(wav_file_path, whisper_model=whisper_model)
251
+
252
+ # Handle segments nested under 'segments' key
253
+ if isinstance(segments, dict) and 'segments' in segments:
254
+ segments = segments['segments']
255
+
256
+ if isinstance(segments, list):
257
+ # Log first 5 segments for debugging
258
+ logging.debug(f"Segments before formatting: {segments[:5]}")
259
+ transcription = format_transcription_with_timestamps(segments)
260
+ logging.debug(f"Formatted transcription (first 500 chars): {transcription[:500]}")
261
+ update_progress("Audio transcribed successfully.")
262
+ else:
263
+ update_progress("Unexpected segments format received from speech_to_text.")
264
+ logging.error(f"Unexpected segments format: {segments}")
265
+ failed_count += 1
266
+ log_counter(
267
+ metric_name="audio_files_failed_total",
268
+ labels={"whisper_model": whisper_model, "api_name": api_name},
269
+ value=1
270
+ )
271
+ continue
272
+
273
+ if not transcription.strip():
274
+ update_progress("Transcription is empty.")
275
+ failed_count += 1
276
+ log_counter(
277
+ metric_name="audio_files_failed_total",
278
+ labels={"whisper_model": whisper_model, "api_name": api_name},
279
+ value=1
280
+ )
281
+ else:
282
+ # Apply chunking
283
+ chunked_text = improved_chunking_process(transcription, chunk_options)
284
+
285
+ # Summarize
286
+ logging.debug(f"Audio Transcription API Name: {api_name}")
287
+ if api_name:
288
+ try:
289
+ summary = perform_summarization(api_name, chunked_text, custom_prompt_input, api_key)
290
+ update_progress("Audio summarized successfully.")
291
+ except Exception as e:
292
+ logging.error(f"Error during summarization: {str(e)}")
293
+ summary = "Summary generation failed"
294
+ failed_count += 1
295
+ log_counter(
296
+ metric_name="audio_files_failed_total",
297
+ labels={"whisper_model": whisper_model, "api_name": api_name},
298
+ value=1
299
+ )
300
+ else:
301
+ summary = "No summary available (API not provided)"
302
+
303
+ all_transcriptions.append(transcription)
304
+ all_summaries.append(summary)
305
+
306
+ # Use custom_title if provided, otherwise use the original filename
307
+ title = custom_title if custom_title else os.path.basename(wav_file_path)
308
+
309
+ # Add to database
310
+ add_media_with_keywords(
311
+ url=url,
312
+ title=title,
313
+ media_type='audio',
314
+ content=transcription,
315
+ keywords=custom_keywords,
316
+ prompt=custom_prompt_input,
317
+ summary=summary,
318
+ transcription_model=whisper_model,
319
+ author="Unknown",
320
+ ingestion_date=datetime.now().strftime('%Y-%m-%d')
321
+ )
322
+ update_progress("Audio file processed and added to database.")
323
+ processed_count += 1
324
+ log_counter(
325
+ metric_name="audio_files_processed_total",
326
+ labels={"whisper_model": whisper_model, "api_name": api_name},
327
+ value=1
328
+ )
329
+
330
+ # Process uploaded file if provided
331
+ if audio_file:
332
+ url = generate_unique_id()
333
+ if os.path.getsize(audio_file.name) > MAX_FILE_SIZE:
334
+ update_progress(
335
+ f"Uploaded file size exceeds the maximum limit of {MAX_FILE_SIZE / (1024 * 1024):.2f}MB. Skipping this file.")
336
+ else:
337
+ try:
338
+ # Re-encode MP3 to fix potential issues
339
+ reencoded_mp3_path = reencode_mp3(audio_file.name)
340
+ if not os.path.exists(reencoded_mp3_path):
341
+ update_progress(f"Re-encoded file not found: {reencoded_mp3_path}")
342
+ return update_progress("Processing failed: Re-encoded file not found"), "", ""
343
+
344
+ temp_files.append(reencoded_mp3_path)
345
+
346
+ # Convert re-encoded MP3 to WAV
347
+ wav_file_path = convert_mp3_to_wav(reencoded_mp3_path)
348
+ if not os.path.exists(wav_file_path):
349
+ update_progress(f"Converted WAV file not found: {wav_file_path}")
350
+ return update_progress("Processing failed: Converted WAV file not found"), "", ""
351
+
352
+ temp_files.append(wav_file_path)
353
+
354
+ # Initialize transcription
355
+ transcription = ""
356
+
357
+ if diarize:
358
+ segments = speech_to_text(wav_file_path, whisper_model=whisper_model, diarize=True)
359
+ else:
360
+ segments = speech_to_text(wav_file_path, whisper_model=whisper_model)
361
+
362
+ # Handle segments nested under 'segments' key
363
+ if isinstance(segments, dict) and 'segments' in segments:
364
+ segments = segments['segments']
365
+
366
+ if isinstance(segments, list):
367
+ transcription = format_transcription_with_timestamps(segments)
368
+ else:
369
+ update_progress("Unexpected segments format received from speech_to_text.")
370
+ logging.error(f"Unexpected segments format: {segments}")
371
+
372
+ chunked_text = improved_chunking_process(transcription, chunk_options)
373
+
374
+ logging.debug(f"Audio Transcription API Name: {api_name}")
375
+ if api_name:
376
+ try:
377
+ summary = perform_summarization(api_name, chunked_text, custom_prompt_input, api_key)
378
+ update_progress("Audio summarized successfully.")
379
+ except Exception as e:
380
+ logging.error(f"Error during summarization: {str(e)}")
381
+ summary = "Summary generation failed"
382
+ else:
383
+ summary = "No summary available (API not provided)"
384
+
385
+ all_transcriptions.append(transcription)
386
+ all_summaries.append(summary)
387
+
388
+ # Use custom_title if provided, otherwise use the original filename
389
+ title = custom_title if custom_title else os.path.basename(wav_file_path)
390
+
391
+ add_media_with_keywords(
392
+ url="Uploaded File",
393
+ title=title,
394
+ media_type='audio',
395
+ content=transcription,
396
+ keywords=custom_keywords,
397
+ prompt=custom_prompt_input,
398
+ summary=summary,
399
+ transcription_model=whisper_model,
400
+ author="Unknown",
401
+ ingestion_date=datetime.now().strftime('%Y-%m-%d')
402
+ )
403
+ update_progress("Uploaded file processed and added to database.")
404
+ processed_count += 1
405
+ log_counter(
406
+ metric_name="audio_files_processed_total",
407
+ labels={"whisper_model": whisper_model, "api_name": api_name},
408
+ value=1
409
+ )
410
+ except Exception as e:
411
+ update_progress(f"Error processing uploaded file: {str(e)}")
412
+ logging.error(f"Error processing uploaded file: {str(e)}")
413
+ failed_count += 1
414
+ log_counter(
415
+ metric_name="audio_files_failed_total",
416
+ labels={"whisper_model": whisper_model, "api_name": api_name},
417
+ value=1
418
+ )
419
+ return update_progress("Processing failed: Error processing uploaded file"), "", ""
420
+ # Final cleanup
421
+ if not keep_original:
422
+ cleanup_files()
423
+
424
+ end_time = time.time()
425
+ processing_time = end_time - start_time
426
+ # Log processing time
427
+ log_histogram(
428
+ metric_name="audio_processing_time_seconds",
429
+ value=processing_time,
430
+ labels={"whisper_model": whisper_model, "api_name": api_name}
431
+ )
432
+
433
+ # Optionally, log total counts
434
+ log_counter(
435
+ metric_name="total_audio_files_processed",
436
+ labels={"whisper_model": whisper_model, "api_name": api_name},
437
+ value=processed_count
438
+ )
439
+
440
+ log_counter(
441
+ metric_name="total_audio_files_failed",
442
+ labels={"whisper_model": whisper_model, "api_name": api_name},
443
+ value=failed_count
444
+ )
445
+
446
+
447
+ final_progress = update_progress("All processing complete.")
448
+ final_transcriptions = "\n\n".join(all_transcriptions)
449
+ final_summaries = "\n\n".join(all_summaries)
450
+
451
+ return final_progress, final_transcriptions, final_summaries
452
+
453
+ except Exception as e:
454
+ logging.error(f"Error processing audio files: {str(e)}")
455
+ log_counter(
456
+ metric_name="audio_files_failed_total",
457
+ labels={"whisper_model": whisper_model, "api_name": api_name},
458
+ value=1
459
+ )
460
+ cleanup_files()
461
+ return update_progress(f"Processing failed: {str(e)}"), "", ""
462
+
463
+
464
+ def format_transcription_with_timestamps(segments, keep_timestamps):
465
+ """
466
+ Formats the transcription segments with or without timestamps.
467
+
468
+ Parameters:
469
+ segments (list): List of transcription segments.
470
+ keep_timestamps (bool): Whether to include timestamps.
471
+
472
+ Returns:
473
+ str: Formatted transcription.
474
+ """
475
+ if keep_timestamps:
476
+ formatted_segments = []
477
+ for segment in segments:
478
+ start = segment.get('Time_Start', 0)
479
+ end = segment.get('Time_End', 0)
480
+ text = segment.get('Text', '').strip()
481
+
482
+ formatted_segments.append(f"[{start:.2f}-{end:.2f}] {text}")
483
+ return "\n".join(formatted_segments)
484
+ else:
485
+ return "\n".join([segment.get('Text', '').strip() for segment in segments])
486
+
487
+
488
+ def download_youtube_audio(url):
489
+ try:
490
+ # Determine ffmpeg path based on the operating system.
491
+ ffmpeg_path = './Bin/ffmpeg.exe' if os.name == 'nt' else 'ffmpeg'
492
+
493
+ # Create a temporary directory
494
+ with tempfile.TemporaryDirectory() as temp_dir:
495
+ # Extract information about the video
496
+ with yt_dlp.YoutubeDL({'quiet': True}) as ydl:
497
+ info_dict = ydl.extract_info(url, download=False)
498
+ sanitized_title = sanitize_filename(info_dict['title'])
499
+
500
+ # Setup the temporary filenames
501
+ temp_video_path = Path(temp_dir) / f"{sanitized_title}_temp.mp4"
502
+ temp_audio_path = Path(temp_dir) / f"{sanitized_title}.mp3"
503
+
504
+ # Initialize yt-dlp with options for downloading
505
+ ydl_opts = {
506
+ 'format': 'bestaudio[ext=m4a]/best[height<=480]', # Prefer best audio, or video up to 480p
507
+ 'ffmpeg_location': ffmpeg_path,
508
+ 'outtmpl': str(temp_video_path),
509
+ 'noplaylist': True,
510
+ 'quiet': True
511
+ }
512
+
513
+ # Execute yt-dlp to download the video/audio
514
+ with yt_dlp.YoutubeDL(ydl_opts) as ydl:
515
+ ydl.download([url])
516
+
517
+ # Check if the file exists
518
+ if not temp_video_path.exists():
519
+ raise FileNotFoundError(f"Expected file was not found: {temp_video_path}")
520
+
521
+ # Use ffmpeg to extract audio
522
+ ffmpeg_command = [
523
+ ffmpeg_path,
524
+ '-i', str(temp_video_path),
525
+ '-vn', # No video
526
+ '-acodec', 'libmp3lame',
527
+ '-b:a', '192k',
528
+ str(temp_audio_path)
529
+ ]
530
+ subprocess.run(ffmpeg_command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
531
+
532
+ # Check if the audio file was created
533
+ if not temp_audio_path.exists():
534
+ raise FileNotFoundError(f"Expected audio file was not found: {temp_audio_path}")
535
+
536
+ # Create a persistent directory for the download if it doesn't exist
537
+ persistent_dir = Path("downloads")
538
+ persistent_dir.mkdir(exist_ok=True)
539
+
540
+ # Move the file from the temporary directory to the persistent directory
541
+ persistent_file_path = persistent_dir / f"{sanitized_title}.mp3"
542
+ os.replace(str(temp_audio_path), str(persistent_file_path))
543
+
544
+ # Add the file to the list of downloaded files
545
+ downloaded_files.append(str(persistent_file_path))
546
+
547
+ return str(persistent_file_path), f"Audio downloaded successfully: {sanitized_title}.mp3"
548
+ except Exception as e:
549
+ return None, f"Error downloading audio: {str(e)}"
550
+
551
+
552
+ def process_podcast(url, title, author, keywords, custom_prompt, api_name, api_key, whisper_model,
553
+ keep_original=False, enable_diarization=False, use_cookies=False, cookies=None,
554
+ chunk_method=None, max_chunk_size=300, chunk_overlap=0, use_adaptive_chunking=False,
555
+ use_multi_level_chunking=False, chunk_language='english', keep_timestamps=True):
556
+ """
557
+ Processes a podcast by downloading the audio, transcribing it, summarizing the transcription,
558
+ and adding the results to the database. Metrics are logged throughout the process.
559
+
560
+ Parameters:
561
+ url (str): URL of the podcast.
562
+ title (str): Title of the podcast.
563
+ author (str): Author of the podcast.
564
+ keywords (str): Comma-separated keywords.
565
+ custom_prompt (str): Custom prompt for summarization.
566
+ api_name (str): API name for summarization.
567
+ api_key (str): API key for summarization.
568
+ whisper_model (str): Whisper model to use for transcription.
569
+ keep_original (bool): Whether to keep the original audio file.
570
+ enable_diarization (bool): Whether to enable speaker diarization.
571
+ use_cookies (bool): Whether to use cookies for authenticated downloads.
572
+ cookies (str): JSON-formatted cookies string.
573
+ chunk_method (str): Method for chunking text.
574
+ max_chunk_size (int): Maximum size for each text chunk.
575
+ chunk_overlap (int): Overlap size between chunks.
576
+ use_adaptive_chunking (bool): Whether to use adaptive chunking.
577
+ use_multi_level_chunking (bool): Whether to use multi-level chunking.
578
+ chunk_language (str): Language for chunking.
579
+ keep_timestamps (bool): Whether to keep timestamps in transcription.
580
+
581
+ Returns:
582
+ tuple: (progress_message, transcription, summary, title, author, keywords, error_message)
583
+ """
584
+ start_time = time.time() # Start time for processing
585
+ error_message = ""
586
+ temp_files = []
587
+
588
+ # Define labels for metrics
589
+ labels = {
590
+ "whisper_model": whisper_model,
591
+ "api_name": api_name if api_name else "None"
592
+ }
593
+
594
+ def update_progress(message):
595
+ """
596
+ Updates the progress messages.
597
+
598
+ Parameters:
599
+ message (str): Progress message to append.
600
+
601
+ Returns:
602
+ str: Combined progress messages.
603
+ """
604
+ progress.append(message)
605
+ return "\n".join(progress)
606
+
607
+ def cleanup_files():
608
+ if not keep_original:
609
+ for file in temp_files:
610
+ try:
611
+ if os.path.exists(file):
612
+ os.remove(file)
613
+ update_progress(f"Temporary file {file} removed.")
614
+ except Exception as e:
615
+ update_progress(f"Failed to remove temporary file {file}: {str(e)}")
616
+
617
+ progress = [] # Initialize progress messages
618
+
619
+ try:
620
+ # Handle cookies if required
621
+ if use_cookies:
622
+ cookies = json.loads(cookies)
623
+
624
+ # Download the podcast audio file
625
+ audio_file = download_audio_file(url, whisper_model, use_cookies, cookies)
626
+ if not audio_file:
627
+ raise RuntimeError("Failed to download podcast audio.")
628
+ temp_files.append(audio_file)
629
+ update_progress("Podcast downloaded successfully.")
630
+
631
+ # Extract metadata from the podcast
632
+ metadata = extract_metadata(url)
633
+ title = title or metadata.get('title', 'Unknown Podcast')
634
+ author = author or metadata.get('uploader', 'Unknown Author')
635
+
636
+ # Format metadata for storage
637
+ metadata_text = f"""
638
+ Metadata:
639
+ Title: {title}
640
+ Author: {author}
641
+ Series: {metadata.get('series', 'N/A')}
642
+ Episode: {metadata.get('episode', 'N/A')}
643
+ Season: {metadata.get('season', 'N/A')}
644
+ Upload Date: {metadata.get('upload_date', 'N/A')}
645
+ Duration: {metadata.get('duration', 'N/A')} seconds
646
+ Description: {metadata.get('description', 'N/A')}
647
+ """
648
+
649
+ # Update keywords with metadata information
650
+ new_keywords = []
651
+ if metadata.get('series'):
652
+ new_keywords.append(f"series:{metadata['series']}")
653
+ if metadata.get('episode'):
654
+ new_keywords.append(f"episode:{metadata['episode']}")
655
+ if metadata.get('season'):
656
+ new_keywords.append(f"season:{metadata['season']}")
657
+
658
+ keywords = f"{keywords},{','.join(new_keywords)}" if keywords else ','.join(new_keywords)
659
+ update_progress(f"Metadata extracted - Title: {title}, Author: {author}, Keywords: {keywords}")
660
+
661
+ # Transcribe the podcast audio
662
+ try:
663
+ if enable_diarization:
664
+ segments = speech_to_text(audio_file, whisper_model=whisper_model, diarize=True)
665
+ else:
666
+ segments = speech_to_text(audio_file, whisper_model=whisper_model)
667
+ # SEems like this could be optimized... FIXME
668
+ def format_segment(segment):
669
+ start = segment.get('start', 0)
670
+ end = segment.get('end', 0)
671
+ text = segment.get('Text', '')
672
+
673
+ if isinstance(segments, dict) and 'segments' in segments:
674
+ segments = segments['segments']
675
+
676
+ if isinstance(segments, list):
677
+ transcription = format_transcription_with_timestamps(segments, keep_timestamps)
678
+ update_progress("Podcast transcribed successfully.")
679
+ else:
680
+ raise ValueError("Unexpected segments format received from speech_to_text.")
681
+
682
+ if not transcription.strip():
683
+ raise ValueError("Transcription is empty.")
684
+ except Exception as e:
685
+ error_message = f"Transcription failed: {str(e)}"
686
+ raise RuntimeError(error_message)
687
+
688
+ # Apply chunking to the transcription
689
+ chunk_options = {
690
+ 'method': chunk_method,
691
+ 'max_size': max_chunk_size,
692
+ 'overlap': chunk_overlap,
693
+ 'adaptive': use_adaptive_chunking,
694
+ 'multi_level': use_multi_level_chunking,
695
+ 'language': chunk_language
696
+ }
697
+ chunked_text = improved_chunking_process(transcription, chunk_options)
698
+
699
+ # Combine metadata and transcription
700
+ full_content = metadata_text + "\n\nTranscription:\n" + transcription
701
+
702
+ # Summarize the transcription if API is provided
703
+ summary = None
704
+ if api_name:
705
+ try:
706
+ summary = perform_summarization(api_name, chunked_text, custom_prompt, api_key)
707
+ update_progress("Podcast summarized successfully.")
708
+ except Exception as e:
709
+ error_message = f"Summarization failed: {str(e)}"
710
+ raise RuntimeError(error_message)
711
+ else:
712
+ summary = "No summary available (API not provided)"
713
+
714
+ # Add the processed podcast to the database
715
+ try:
716
+ add_media_with_keywords(
717
+ url=url,
718
+ title=title,
719
+ media_type='podcast',
720
+ content=full_content,
721
+ keywords=keywords,
722
+ prompt=custom_prompt,
723
+ summary=summary or "No summary available",
724
+ transcription_model=whisper_model,
725
+ author=author,
726
+ ingestion_date=datetime.now().strftime('%Y-%m-%d')
727
+ )
728
+ update_progress("Podcast added to database successfully.")
729
+ except Exception as e:
730
+ error_message = f"Error adding podcast to database: {str(e)}"
731
+ raise RuntimeError(error_message)
732
+
733
+ # Cleanup temporary files if required
734
+ cleanup_files()
735
+
736
+ # Calculate processing time
737
+ end_time = time.time()
738
+ processing_time = end_time - start_time
739
+
740
+ # Log successful processing
741
+ log_counter(
742
+ metric_name="podcasts_processed_total",
743
+ labels=labels,
744
+ value=1
745
+ )
746
+
747
+ # Log processing time
748
+ log_histogram(
749
+ metric_name="podcast_processing_time_seconds",
750
+ value=processing_time,
751
+ labels=labels
752
+ )
753
+
754
+ # Return the final outputs
755
+ final_progress = update_progress("Processing complete.")
756
+ return (final_progress, full_content, summary or "No summary generated.",
757
+ title, author, keywords, error_message)
758
+
759
+ except Exception as e:
760
+ # Calculate processing time up to the point of failure
761
+ end_time = time.time()
762
+ processing_time = end_time - start_time
763
+
764
+ # Log failed processing
765
+ log_counter(
766
+ metric_name="podcasts_failed_total",
767
+ labels=labels,
768
+ value=1
769
+ )
770
+
771
+ # Log processing time even on failure
772
+ log_histogram(
773
+ metric_name="podcast_processing_time_seconds",
774
+ value=processing_time,
775
+ labels=labels
776
+ )
777
+
778
+ logging.error(f"Error processing podcast: {str(e)}")
779
+ cleanup_files()
780
+ final_progress = update_progress(f"Processing failed: {str(e)}")
781
+ return (final_progress, "", "", "", "", "", str(e))
782
+
783
+
784
+ #
785
+ #
786
+ #######################################################################################################################
App_Function_Libraries/Audio/Audio_Transcription_Lib.py ADDED
@@ -0,0 +1,335 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Audio_Transcription_Lib.py
2
+ #########################################
3
+ # Transcription Library
4
+ # This library is used to perform transcription of audio files.
5
+ # Currently, uses faster_whisper for transcription.
6
+ #
7
+ ####################
8
+ # Function List
9
+ #
10
+ # 1. convert_to_wav(video_file_path, offset=0, overwrite=False)
11
+ # 2. speech_to_text(audio_file_path, selected_source_lang='en', whisper_model='small.en', vad_filter=False)
12
+ #
13
+ ####################
14
+ #
15
+ # Import necessary libraries to run solo for testing
16
+ import gc
17
+ import json
18
+ import logging
19
+ import multiprocessing
20
+ import os
21
+ import queue
22
+ import sys
23
+ import subprocess
24
+ import tempfile
25
+ import threading
26
+ import time
27
+ # DEBUG Imports
28
+ #from memory_profiler import profile
29
+ import pyaudio
30
+ from faster_whisper import WhisperModel as OriginalWhisperModel
31
+ from typing import Optional, Union, List, Dict, Any
32
+ #
33
+ # Import Local
34
+ from App_Function_Libraries.Utils.Utils import load_comprehensive_config
35
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
36
+ #
37
+ #######################################################################################################################
38
+ # Function Definitions
39
+ #
40
+
41
+ # Convert video .m4a into .wav using ffmpeg
42
+ # ffmpeg -i "example.mp4" -ar 16000 -ac 1 -c:a pcm_s16le "output.wav"
43
+ # https://www.gyan.dev/ffmpeg/builds/
44
+ #
45
+
46
+
47
+ whisper_model_instance = None
48
+ config = load_comprehensive_config()
49
+ processing_choice = config.get('Processing', 'processing_choice', fallback='cpu')
50
+ total_thread_count = multiprocessing.cpu_count()
51
+
52
+
53
+ class WhisperModel(OriginalWhisperModel):
54
+ tldw_dir = os.path.dirname(os.path.dirname(__file__))
55
+ default_download_root = os.path.join(tldw_dir, 'models', 'Whisper')
56
+
57
+ valid_model_sizes = [
58
+ "tiny.en", "tiny", "base.en", "base", "small.en", "small", "medium.en", "medium",
59
+ "large-v1", "large-v2", "large-v3", "large", "distil-large-v2", "distil-medium.en",
60
+ "distil-small.en", "distil-large-v3",
61
+ ]
62
+
63
+ def __init__(
64
+ self,
65
+ model_size_or_path: str,
66
+ device: str = processing_choice,
67
+ device_index: Union[int, List[int]] = 0,
68
+ compute_type: str = "default",
69
+ cpu_threads: int = 0,#total_thread_count, FIXME - I think this should be 0
70
+ num_workers: int = 1,
71
+ download_root: Optional[str] = None,
72
+ local_files_only: bool = False,
73
+ files: Optional[Dict[str, Any]] = None,
74
+ **model_kwargs: Any
75
+ ):
76
+ if download_root is None:
77
+ download_root = self.default_download_root
78
+
79
+ os.makedirs(download_root, exist_ok=True)
80
+
81
+ # FIXME - validate....
82
+ # Also write an integration test...
83
+ # Check if model_size_or_path is a valid model size
84
+ if model_size_or_path in self.valid_model_sizes:
85
+ # It's a model size, so we'll use the download_root
86
+ model_path = os.path.join(download_root, model_size_or_path)
87
+ if not os.path.isdir(model_path):
88
+ # If it doesn't exist, we'll let the parent class download it
89
+ model_size_or_path = model_size_or_path # Keep the original model size
90
+ else:
91
+ # If it exists, use the full path
92
+ model_size_or_path = model_path
93
+ else:
94
+ # It's not a valid model size, so assume it's a path
95
+ model_size_or_path = os.path.abspath(model_size_or_path)
96
+
97
+ super().__init__(
98
+ model_size_or_path,
99
+ device=device,
100
+ device_index=device_index,
101
+ compute_type=compute_type,
102
+ cpu_threads=cpu_threads,
103
+ num_workers=num_workers,
104
+ download_root=download_root,
105
+ local_files_only=local_files_only,
106
+ # Maybe? idk, FIXME
107
+ # files=files,
108
+ # **model_kwargs
109
+ )
110
+
111
+ def get_whisper_model(model_name, device):
112
+ global whisper_model_instance
113
+ if whisper_model_instance is None:
114
+ logging.info(f"Initializing new WhisperModel with size {model_name} on device {device}")
115
+ whisper_model_instance = WhisperModel(model_name, device=device)
116
+ return whisper_model_instance
117
+
118
+ # os.system(r'.\Bin\ffmpeg.exe -ss 00:00:00 -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{out_path}"')
119
+ #DEBUG
120
+ #@profile
121
+ def convert_to_wav(video_file_path, offset=0, overwrite=False):
122
+ log_counter("convert_to_wav_attempt", labels={"file_path": video_file_path})
123
+ start_time = time.time()
124
+
125
+ out_path = os.path.splitext(video_file_path)[0] + ".wav"
126
+
127
+ if os.path.exists(out_path) and not overwrite:
128
+ print(f"File '{out_path}' already exists. Skipping conversion.")
129
+ logging.info(f"Skipping conversion as file already exists: {out_path}")
130
+ log_counter("convert_to_wav_skipped", labels={"file_path": video_file_path})
131
+ return out_path
132
+
133
+ print("Starting conversion process of .m4a to .WAV")
134
+ out_path = os.path.splitext(video_file_path)[0] + ".wav"
135
+
136
+ try:
137
+ if os.name == "nt":
138
+ logging.debug("ffmpeg being ran on windows")
139
+
140
+ if sys.platform.startswith('win'):
141
+ ffmpeg_cmd = ".\\Bin\\ffmpeg.exe"
142
+ logging.debug(f"ffmpeg_cmd: {ffmpeg_cmd}")
143
+ else:
144
+ ffmpeg_cmd = 'ffmpeg' # Assume 'ffmpeg' is in PATH for non-Windows systems
145
+
146
+ command = [
147
+ ffmpeg_cmd, # Assuming the working directory is correctly set where .\Bin exists
148
+ "-ss", "00:00:00", # Start at the beginning of the video
149
+ "-i", video_file_path,
150
+ "-ar", "16000", # Audio sample rate
151
+ "-ac", "1", # Number of audio channels
152
+ "-c:a", "pcm_s16le", # Audio codec
153
+ out_path
154
+ ]
155
+ try:
156
+ # Redirect stdin from null device to prevent ffmpeg from waiting for input
157
+ with open(os.devnull, 'rb') as null_file:
158
+ result = subprocess.run(command, stdin=null_file, text=True, capture_output=True)
159
+ if result.returncode == 0:
160
+ logging.info("FFmpeg executed successfully")
161
+ logging.debug("FFmpeg output: %s", result.stdout)
162
+ else:
163
+ logging.error("Error in running FFmpeg")
164
+ logging.error("FFmpeg stderr: %s", result.stderr)
165
+ raise RuntimeError(f"FFmpeg error: {result.stderr}")
166
+ except Exception as e:
167
+ logging.error("Error occurred - ffmpeg doesn't like windows")
168
+ raise RuntimeError("ffmpeg failed")
169
+ elif os.name == "posix":
170
+ os.system(f'ffmpeg -ss 00:00:00 -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{out_path}"')
171
+ else:
172
+ raise RuntimeError("Unsupported operating system")
173
+ logging.info("Conversion to WAV completed: %s", out_path)
174
+ log_counter("convert_to_wav_success", labels={"file_path": video_file_path})
175
+ except Exception as e:
176
+ logging.error("speech-to-text: Error transcribing audio: %s", str(e))
177
+ log_counter("convert_to_wav_error", labels={"file_path": video_file_path, "error": str(e)})
178
+ return {"error": str(e)}
179
+
180
+ conversion_time = time.time() - start_time
181
+ log_histogram("convert_to_wav_duration", conversion_time, labels={"file_path": video_file_path})
182
+
183
+ gc.collect()
184
+ return out_path
185
+
186
+
187
+ # Transcribe .wav into .segments.json
188
+ #DEBUG
189
+ #@profile
190
+ # FIXME - I feel like the `vad_filter` shoudl be enabled by default....
191
+ def speech_to_text(audio_file_path, selected_source_lang='en', whisper_model='medium.en', vad_filter=False, diarize=False):
192
+ log_counter("speech_to_text_attempt", labels={"file_path": audio_file_path, "model": whisper_model})
193
+ time_start = time.time()
194
+
195
+ if audio_file_path is None:
196
+ log_counter("speech_to_text_error", labels={"error": "No audio file provided"})
197
+ raise ValueError("speech-to-text: No audio file provided")
198
+ logging.info("speech-to-text: Audio file path: %s", audio_file_path)
199
+
200
+ try:
201
+ _, file_ending = os.path.splitext(audio_file_path)
202
+ out_file = audio_file_path.replace(file_ending, "-whisper_model-"+whisper_model+".segments.json")
203
+ prettified_out_file = audio_file_path.replace(file_ending, "-whisper_model-"+whisper_model+".segments_pretty.json")
204
+ if os.path.exists(out_file):
205
+ logging.info("speech-to-text: Segments file already exists: %s", out_file)
206
+ with open(out_file) as f:
207
+ global segments
208
+ segments = json.load(f)
209
+ return segments
210
+
211
+ logging.info('speech-to-text: Starting transcription...')
212
+ # FIXME - revisit this
213
+ options = dict(language=selected_source_lang, beam_size=10, best_of=10, vad_filter=vad_filter)
214
+ transcribe_options = dict(task="transcribe", **options)
215
+ # use function and config at top of file
216
+ logging.debug("speech-to-text: Using whisper model: %s", whisper_model)
217
+ whisper_model_instance = get_whisper_model(whisper_model, processing_choice)
218
+ # faster_whisper transcription right here - FIXME -test batching - ha
219
+ segments_raw, info = whisper_model_instance.transcribe(audio_file_path, **transcribe_options)
220
+
221
+ segments = []
222
+ for segment_chunk in segments_raw:
223
+ chunk = {
224
+ "Time_Start": segment_chunk.start,
225
+ "Time_End": segment_chunk.end,
226
+ "Text": segment_chunk.text
227
+ }
228
+ logging.debug("Segment: %s", chunk)
229
+ segments.append(chunk)
230
+ # Print to verify its working
231
+ logging.info(f"{segment_chunk.start:.2f}s - {segment_chunk.end:.2f}s | {segment_chunk.text}")
232
+
233
+ # Log it as well.
234
+ logging.debug(
235
+ f"Transcribed Segment: {segment_chunk.start:.2f}s - {segment_chunk.end:.2f}s | {segment_chunk.text}")
236
+
237
+ if segments:
238
+ segments[0]["Text"] = f"This text was transcribed using whisper model: {whisper_model}\n\n" + segments[0]["Text"]
239
+
240
+ if not segments:
241
+ log_counter("speech_to_text_error", labels={"error": "No transcription produced"})
242
+ raise RuntimeError("No transcription produced. The audio file may be invalid or empty.")
243
+
244
+ transcription_time = time.time() - time_start
245
+ logging.info("speech-to-text: Transcription completed in %.2f seconds", transcription_time)
246
+ log_histogram("speech_to_text_duration", transcription_time, labels={"file_path": audio_file_path, "model": whisper_model})
247
+ log_counter("speech_to_text_success", labels={"file_path": audio_file_path, "model": whisper_model})
248
+ # Save the segments to a JSON file - prettified and non-prettified
249
+ # FIXME refactor so this is an optional flag to save either the prettified json file or the normal one
250
+ save_json = True
251
+ if save_json:
252
+ logging.info("speech-to-text: Saving segments to JSON file")
253
+ output_data = {'segments': segments}
254
+ logging.info("speech-to-text: Saving prettified JSON to %s", prettified_out_file)
255
+ with open(prettified_out_file, 'w') as f:
256
+ json.dump(output_data, f, indent=2)
257
+
258
+ logging.info("speech-to-text: Saving JSON to %s", out_file)
259
+ with open(out_file, 'w') as f:
260
+ json.dump(output_data, f)
261
+
262
+ logging.debug(f"speech-to-text: returning {segments[:500]}")
263
+ gc.collect()
264
+ return segments
265
+
266
+ except Exception as e:
267
+ logging.error("speech-to-text: Error transcribing audio: %s", str(e))
268
+ log_counter("speech_to_text_error", labels={"file_path": audio_file_path, "model": whisper_model, "error": str(e)})
269
+ raise RuntimeError("speech-to-text: Error transcribing audio")
270
+
271
+
272
+ def record_audio(duration, sample_rate=16000, chunk_size=1024):
273
+ log_counter("record_audio_attempt", labels={"duration": duration})
274
+ p = pyaudio.PyAudio()
275
+ stream = p.open(format=pyaudio.paInt16,
276
+ channels=1,
277
+ rate=sample_rate,
278
+ input=True,
279
+ frames_per_buffer=chunk_size)
280
+
281
+ print("Recording...")
282
+ frames = []
283
+ stop_recording = threading.Event()
284
+ audio_queue = queue.Queue()
285
+
286
+ def audio_callback():
287
+ for _ in range(0, int(sample_rate / chunk_size * duration)):
288
+ if stop_recording.is_set():
289
+ break
290
+ data = stream.read(chunk_size)
291
+ audio_queue.put(data)
292
+
293
+ audio_thread = threading.Thread(target=audio_callback)
294
+ audio_thread.start()
295
+
296
+ return p, stream, audio_queue, stop_recording, audio_thread
297
+
298
+
299
+ def stop_recording(p, stream, audio_queue, stop_recording_event, audio_thread):
300
+ log_counter("stop_recording_attempt")
301
+ start_time = time.time()
302
+ stop_recording_event.set()
303
+ audio_thread.join()
304
+
305
+ frames = []
306
+ while not audio_queue.empty():
307
+ frames.append(audio_queue.get())
308
+
309
+ print("Recording finished.")
310
+
311
+ stream.stop_stream()
312
+ stream.close()
313
+ p.terminate()
314
+
315
+ stop_time = time.time() - start_time
316
+ log_histogram("stop_recording_duration", stop_time)
317
+ log_counter("stop_recording_success")
318
+ return b''.join(frames)
319
+
320
+ def save_audio_temp(audio_data, sample_rate=16000):
321
+ log_counter("save_audio_temp_attempt")
322
+ with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
323
+ import wave
324
+ wf = wave.open(temp_file.name, 'wb')
325
+ wf.setnchannels(1)
326
+ wf.setsampwidth(2)
327
+ wf.setframerate(sample_rate)
328
+ wf.writeframes(audio_data)
329
+ wf.close()
330
+ log_counter("save_audio_temp_success")
331
+ return temp_file.name
332
+
333
+ #
334
+ #
335
+ #######################################################################################################################
App_Function_Libraries/Audio/Diarization_Lib.py ADDED
@@ -0,0 +1,275 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Diarization_Lib.py
2
+ #########################################
3
+ # Diarization Library
4
+ # This library is used to perform diarization of audio files.
5
+ # Currently, uses FIXME for transcription.
6
+ #
7
+ ####################
8
+ ####################
9
+ # Function List
10
+ #
11
+ # 1. speaker_diarize(video_file_path, segments, embedding_model = "pyannote/embedding", embedding_size=512, num_speakers=0)
12
+ #
13
+ ####################
14
+ # Import necessary libraries
15
+ import logging
16
+ from pathlib import Path
17
+ from typing import Dict, List, Any
18
+
19
+ #
20
+ # Import Local Libraries
21
+ from App_Function_Libraries.Audio.Audio_Transcription_Lib import speech_to_text
22
+ #
23
+ # Import 3rd Party Libraries
24
+ from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
25
+ import yaml
26
+ #
27
+ #######################################################################################################################
28
+ # Function Definitions
29
+ #
30
+
31
+ def load_pipeline_from_pretrained(path_to_config: str | Path) -> SpeakerDiarization:
32
+ path_to_config = Path(path_to_config).resolve()
33
+ logging.debug(f"Loading pyannote pipeline from {path_to_config}...")
34
+
35
+ if not path_to_config.exists():
36
+ raise FileNotFoundError(f"Config file not found: {path_to_config}")
37
+
38
+ # Load the YAML configuration
39
+ with open(path_to_config, 'r') as config_file:
40
+ config = yaml.safe_load(config_file)
41
+
42
+ # Debug: print the entire config
43
+ logging.debug(f"Loaded config: {config}")
44
+
45
+ # Create the SpeakerDiarization pipeline
46
+ try:
47
+ pipeline = SpeakerDiarization(
48
+ segmentation=config['pipeline']['params']['segmentation'],
49
+ embedding=config['pipeline']['params']['embedding'],
50
+ clustering=config['pipeline']['params']['clustering'],
51
+ )
52
+ except KeyError as e:
53
+ logging.error(f"Error accessing config key: {e}")
54
+ raise
55
+
56
+ # Set other parameters
57
+ try:
58
+ pipeline_params = {
59
+ "segmentation": {},
60
+ "clustering": {},
61
+ }
62
+
63
+ if 'params' in config and 'segmentation' in config['params']:
64
+ if 'min_duration_off' in config['params']['segmentation']:
65
+ pipeline_params["segmentation"]["min_duration_off"] = config['params']['segmentation']['min_duration_off']
66
+
67
+ if 'params' in config and 'clustering' in config['params']:
68
+ if 'method' in config['params']['clustering']:
69
+ pipeline_params["clustering"]["method"] = config['params']['clustering']['method']
70
+ if 'min_cluster_size' in config['params']['clustering']:
71
+ pipeline_params["clustering"]["min_cluster_size"] = config['params']['clustering']['min_cluster_size']
72
+ if 'threshold' in config['params']['clustering']:
73
+ pipeline_params["clustering"]["threshold"] = config['params']['clustering']['threshold']
74
+
75
+ if 'pipeline' in config and 'params' in config['pipeline']:
76
+ if 'embedding_batch_size' in config['pipeline']['params']:
77
+ pipeline_params["embedding_batch_size"] = config['pipeline']['params']['embedding_batch_size']
78
+ if 'embedding_exclude_overlap' in config['pipeline']['params']:
79
+ pipeline_params["embedding_exclude_overlap"] = config['pipeline']['params']['embedding_exclude_overlap']
80
+ if 'segmentation_batch_size' in config['pipeline']['params']:
81
+ pipeline_params["segmentation_batch_size"] = config['pipeline']['params']['segmentation_batch_size']
82
+
83
+ logging.debug(f"Pipeline params: {pipeline_params}")
84
+ pipeline.instantiate(pipeline_params)
85
+ except KeyError as e:
86
+ logging.error(f"Error accessing config key: {e}")
87
+ raise
88
+ except Exception as e:
89
+ logging.error(f"Error instantiating pipeline: {e}")
90
+ raise
91
+
92
+ return pipeline
93
+
94
+
95
+ def audio_diarization(audio_file_path: str) -> list:
96
+ logging.info('audio-diarization: Loading pyannote pipeline')
97
+
98
+ base_dir = Path(__file__).parent.resolve()
99
+ config_path = base_dir / 'models' / 'pyannote_diarization_config.yaml'
100
+ logging.info(f"audio-diarization: Loading pipeline from {config_path}")
101
+
102
+ try:
103
+ pipeline = load_pipeline_from_pretrained(config_path)
104
+ except Exception as e:
105
+ logging.error(f"Failed to load pipeline: {str(e)}")
106
+ raise
107
+
108
+ logging.info(f"audio-diarization: Audio file path: {audio_file_path}")
109
+
110
+ try:
111
+ logging.info('audio-diarization: Starting diarization...')
112
+ diarization_result = pipeline(audio_file_path)
113
+
114
+ segments = []
115
+ for turn, _, speaker in diarization_result.itertracks(yield_label=True):
116
+ segment = {
117
+ "start": turn.start,
118
+ "end": turn.end,
119
+ "speaker": speaker
120
+ }
121
+ logging.debug(f"Segment: {segment}")
122
+ segments.append(segment)
123
+ logging.info("audio-diarization: Diarization completed with pyannote")
124
+
125
+ return segments
126
+
127
+ except Exception as e:
128
+ logging.error(f"audio-diarization: Error performing diarization: {str(e)}")
129
+ raise RuntimeError("audio-diarization: Error performing diarization") from e
130
+
131
+
132
+ # Old
133
+ # def audio_diarization(audio_file_path):
134
+ # logging.info('audio-diarization: Loading pyannote pipeline')
135
+ #
136
+ # #config file loading
137
+ # current_dir = os.path.dirname(os.path.abspath(__file__))
138
+ # # Construct the path to the config file
139
+ # config_path = os.path.join(current_dir, 'Config_Files', 'config.txt')
140
+ # # Read the config file
141
+ # config = configparser.ConfigParser()
142
+ # config.read(config_path)
143
+ # processing_choice = config.get('Processing', 'processing_choice', fallback='cpu')
144
+ #
145
+ # base_dir = Path(__file__).parent.resolve()
146
+ # config_path = base_dir / 'models' / 'config.yaml'
147
+ # pipeline = load_pipeline_from_pretrained(config_path)
148
+ #
149
+ # time_start = time.time()
150
+ # if audio_file_path is None:
151
+ # raise ValueError("audio-diarization: No audio file provided")
152
+ # logging.info("audio-diarization: Audio file path: %s", audio_file_path)
153
+ #
154
+ # try:
155
+ # _, file_ending = os.path.splitext(audio_file_path)
156
+ # out_file = audio_file_path.replace(file_ending, ".diarization.json")
157
+ # prettified_out_file = audio_file_path.replace(file_ending, ".diarization_pretty.json")
158
+ # if os.path.exists(out_file):
159
+ # logging.info("audio-diarization: Diarization file already exists: %s", out_file)
160
+ # with open(out_file) as f:
161
+ # global diarization_result
162
+ # diarization_result = json.load(f)
163
+ # return diarization_result
164
+ #
165
+ # logging.info('audio-diarization: Starting diarization...')
166
+ # diarization_result = pipeline(audio_file_path)
167
+ #
168
+ # segments = []
169
+ # for turn, _, speaker in diarization_result.itertracks(yield_label=True):
170
+ # chunk = {
171
+ # "Time_Start": turn.start,
172
+ # "Time_End": turn.end,
173
+ # "Speaker": speaker
174
+ # }
175
+ # logging.debug("Segment: %s", chunk)
176
+ # segments.append(chunk)
177
+ # logging.info("audio-diarization: Diarization completed with pyannote")
178
+ #
179
+ # output_data = {'segments': segments}
180
+ #
181
+ # logging.info("audio-diarization: Saving prettified JSON to %s", prettified_out_file)
182
+ # with open(prettified_out_file, 'w') as f:
183
+ # json.dump(output_data, f, indent=2)
184
+ #
185
+ # logging.info("audio-diarization: Saving JSON to %s", out_file)
186
+ # with open(out_file, 'w') as f:
187
+ # json.dump(output_data, f)
188
+ #
189
+ # except Exception as e:
190
+ # logging.error("audio-diarization: Error performing diarization: %s", str(e))
191
+ # raise RuntimeError("audio-diarization: Error performing diarization")
192
+ # return segments
193
+
194
+ def combine_transcription_and_diarization(audio_file_path: str) -> List[Dict[str, Any]]:
195
+ logging.info('combine-transcription-and-diarization: Starting transcription and diarization...')
196
+
197
+ try:
198
+ logging.info('Performing speech-to-text...')
199
+ transcription_result = speech_to_text(audio_file_path)
200
+ logging.info(f"Transcription result type: {type(transcription_result)}")
201
+ logging.info(f"Transcription result: {transcription_result[:3] if isinstance(transcription_result, list) and len(transcription_result) > 3 else transcription_result}")
202
+
203
+ logging.info('Performing audio diarization...')
204
+ diarization_result = audio_diarization(audio_file_path)
205
+ logging.info(f"Diarization result type: {type(diarization_result)}")
206
+ logging.info(f"Diarization result sample: {diarization_result[:3] if isinstance(diarization_result, list) and len(diarization_result) > 3 else diarization_result}")
207
+
208
+ if not transcription_result:
209
+ logging.error("Empty result from transcription")
210
+ return []
211
+
212
+ if not diarization_result:
213
+ logging.error("Empty result from diarization")
214
+ return []
215
+
216
+ # Handle the case where transcription_result is a dict with a 'segments' key
217
+ if isinstance(transcription_result, dict) and 'segments' in transcription_result:
218
+ transcription_segments = transcription_result['segments']
219
+ elif isinstance(transcription_result, list):
220
+ transcription_segments = transcription_result
221
+ else:
222
+ logging.error(f"Unexpected transcription result format: {type(transcription_result)}")
223
+ return []
224
+
225
+ logging.info(f"Number of transcription segments: {len(transcription_segments)}")
226
+ logging.info(f"Transcription segments sample: {transcription_segments[:3] if len(transcription_segments) > 3 else transcription_segments}")
227
+
228
+ if not isinstance(diarization_result, list):
229
+ logging.error(f"Unexpected diarization result format: {type(diarization_result)}")
230
+ return []
231
+
232
+ combined_result = []
233
+ for transcription_segment in transcription_segments:
234
+ if not isinstance(transcription_segment, dict):
235
+ logging.warning(f"Unexpected transcription segment format: {transcription_segment}")
236
+ continue
237
+
238
+ for diarization_segment in diarization_result:
239
+ if not isinstance(diarization_segment, dict):
240
+ logging.warning(f"Unexpected diarization segment format: {diarization_segment}")
241
+ continue
242
+
243
+ try:
244
+ trans_start = transcription_segment.get('Time_Start', 0)
245
+ trans_end = transcription_segment.get('Time_End', 0)
246
+ diar_start = diarization_segment.get('start', 0)
247
+ diar_end = diarization_segment.get('end', 0)
248
+
249
+ if trans_start >= diar_start and trans_end <= diar_end:
250
+ combined_segment = {
251
+ "Time_Start": trans_start,
252
+ "Time_End": trans_end,
253
+ "Speaker": diarization_segment.get('speaker', 'Unknown'),
254
+ "Text": transcription_segment.get('Text', '')
255
+ }
256
+ combined_result.append(combined_segment)
257
+ break
258
+ except Exception as e:
259
+ logging.error(f"Error processing segment: {str(e)}")
260
+ logging.error(f"Transcription segment: {transcription_segment}")
261
+ logging.error(f"Diarization segment: {diarization_segment}")
262
+ continue
263
+
264
+ logging.info(f"Combined result length: {len(combined_result)}")
265
+ logging.info(f"Combined result sample: {combined_result[:3] if len(combined_result) > 3 else combined_result}")
266
+ return combined_result
267
+
268
+ except Exception as e:
269
+ logging.error(f"Error in combine_transcription_and_diarization: {str(e)}", exc_info=True)
270
+ return []
271
+
272
+
273
+ #
274
+ #
275
+ #######################################################################################################################
App_Function_Libraries/Audio/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/Confabulation_check.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Confabulation_check.py
2
+ #
3
+ # This file contains the functions that are used to check the confabulation of the user's input.
4
+ #
5
+ #
6
+ # Imports
7
+ #
8
+ # External Imports
9
+ #
10
+ # Local Imports
11
+ #
12
+ #
13
+ ####################################################################################################
14
+ #
15
+ # Functions:
16
+ from App_Function_Libraries.Chat import chat_api_call
17
+ from App_Function_Libraries.Benchmarks_Evaluations.ms_g_eval import validate_inputs, detailed_api_error
18
+
19
+
20
+ def simplified_geval(transcript: str, summary: str, api_name: str, api_key: str, temp: float = 0.7) -> str:
21
+ """
22
+ Perform a simplified version of G-Eval using a single query to evaluate the summary.
23
+
24
+ Args:
25
+ transcript (str): The original transcript
26
+ summary (str): The summary to be evaluated
27
+ api_name (str): The name of the LLM API to use
28
+ api_key (str): The API key for the chosen LLM
29
+ temp (float, optional): The temperature parameter for the API call. Defaults to 0.7.
30
+
31
+ Returns:
32
+ str: The evaluation result
33
+ """
34
+ try:
35
+ validate_inputs(transcript, summary, api_name, api_key)
36
+ except ValueError as e:
37
+ return str(e)
38
+
39
+ prompt = f"""You are an AI assistant tasked with evaluating the quality of a summary. You will be given an original transcript and a summary of that transcript. Your task is to evaluate the summary based on the following criteria:
40
+
41
+ 1. Coherence (1-5): How well-structured and organized is the summary?
42
+ 2. Consistency (1-5): How factually aligned is the summary with the original transcript?
43
+ 3. Fluency (1-3): How well-written is the summary in terms of grammar, spelling, and readability?
44
+ 4. Relevance (1-5): How well does the summary capture the important information from the transcript?
45
+
46
+ Please provide a score for each criterion and a brief explanation for your scoring. Then, give an overall assessment of the summary's quality.
47
+
48
+ Original Transcript:
49
+ {transcript}
50
+
51
+ Summary to Evaluate:
52
+ {summary}
53
+
54
+ Please provide your evaluation in the following format:
55
+ Coherence: [score] - [brief explanation]
56
+ Consistency: [score] - [brief explanation]
57
+ Fluency: [score] - [brief explanation]
58
+ Relevance: [score] - [brief explanation]
59
+
60
+ Overall Assessment: [Your overall assessment of the summary's quality]
61
+ """
62
+
63
+ try:
64
+ result = chat_api_call(
65
+ api_name,
66
+ api_key,
67
+ prompt,
68
+ "",
69
+ temp=temp,
70
+ system_message="You are a helpful AI assistant tasked with evaluating summaries."
71
+ )
72
+ except Exception as e:
73
+ return detailed_api_error(api_name, e)
74
+
75
+ formatted_result = f"""
76
+ Confabulation Check Results:
77
+
78
+ {result}
79
+ """
80
+
81
+ return formatted_result
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/.gitignore ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ __pycache__
2
+ .vscode
3
+ *.DS_Store
4
+ *.pyc
5
+ src/plot
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/LICENSE ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ MIT License
2
+
3
+ Copyright (c) 2023 OpenBMB
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
22
+
23
+ taken from https://github.com/OpenBMB/InfiniteBench
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/config.txt ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [API]
2
+ anthropic_api_key = <anthropic_api_key>
3
+ anthropic_model = claude-3-sonnet-20240229
4
+ cohere_api_key = <your_cohere_api_key>
5
+ cohere_model = command-r-plus
6
+ groq_api_key = <your_groq_api_key>
7
+ groq_model = llama3-70b-8192
8
+ openai_api_key = <openai_api_key>
9
+ openai_model = gpt-4-turbo
10
+ huggingface_api_token = <huggingface_api_token>
11
+ huggingface_model = CohereForAI/c4ai-command-r-plus
12
+ openrouter_api_key = <openrouter_api_key>
13
+ openrouter_model = mistralai/mistral-7b-instruct:free
14
+ deepseek_api_key = <deepseek_api_key>
15
+ deepseek_model = deepseek-chat
16
+
17
+ [Local-API]
18
+ kobold_api_key = <kobold api key>
19
+ kobold_api_IP = http://127.0.0.1:5001/api/v1/generate
20
+ llama_api_key = <llama.cpp api key>
21
+ llama_api_IP = http://127.0.0.1:8080/completion
22
+ ooba_api_key = <ooba api key>
23
+ ooba_api_IP = http://127.0.0.1:5000/v1/chat/completions
24
+ tabby_api_IP = http://127.0.0.1:5000/v1/chat/completions
25
+ tabby_api_key = <tabbyapi key>
26
+ vllm_api_IP = http://127.0.0.1:8000/v1/chat/completions
27
+ vllm_model = <vllm model>
28
+ ollama_api_IP = http://127.0.0.1:11434/api/generate
29
+ ollama_api_key = <ollama api key>
30
+ ollama_model = <ollama model>
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_multi_api.py ADDED
@@ -0,0 +1,300 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # eval_multi_api.py
2
+ # Description: Evaluate a language model on a conversational task using multiple APIs
3
+ #
4
+ # Usage: python eval_multi_api.py --task question_answering --api <api_name>> --output_dir ./results --data_dir ./data --verbose
5
+ # API endpoints are defined in the config file (config.txt)
6
+ # The API key for the selected API should be defined in the config file
7
+ # APIs Supported are:
8
+ # - openai
9
+ # - anthropic
10
+ # - cohere
11
+ # - groq
12
+ # - openrouter
13
+ # - deepseek
14
+ # - mistral
15
+ # - llamacpp
16
+ # - kobold
17
+ # - oobabooga
18
+ # - vllm
19
+ # - tabbyapi
20
+ #
21
+ # Imports:
22
+ import configparser
23
+ from pathlib import Path
24
+ import time
25
+ from typing import Dict, Any, Optional, List
26
+ #
27
+ # Local Imports
28
+ from eval_utils import (
29
+ create_msgs,
30
+ load_data,
31
+ dump_jsonl,
32
+ iter_jsonl,
33
+ get_answer,
34
+ )
35
+ from LLM_API_Calls import (
36
+ chat_with_openai,
37
+ chat_with_anthropic,
38
+ chat_with_cohere,
39
+ chat_with_groq,
40
+ chat_with_openrouter,
41
+ chat_with_deepseek,
42
+ chat_with_mistral
43
+ )
44
+ from LLM_API_Calls_Local import (
45
+ chat_with_llama,
46
+ chat_with_kobold,
47
+ chat_with_oobabooga,
48
+ chat_with_vllm,
49
+ chat_with_tabbyapi
50
+ )
51
+ #
52
+ #######################################################################################################################
53
+ #
54
+ # Functions:
55
+
56
+ class MultiAPILLMClient:
57
+ def __init__(self, config_path: str):
58
+ self.config = self.load_config(config_path)
59
+ self.api_functions = {
60
+ 'openai': chat_with_openai,
61
+ 'anthropic': chat_with_anthropic,
62
+ 'cohere': chat_with_cohere,
63
+ 'groq': chat_with_groq,
64
+ 'openrouter': chat_with_openrouter,
65
+ 'deepseek': chat_with_deepseek,
66
+ 'mistral': chat_with_mistral,
67
+ 'llamacpp': chat_with_llama,
68
+ 'kobold': chat_with_kobold,
69
+ 'oobabooga': chat_with_oobabooga,
70
+ 'vllm': chat_with_vllm,
71
+ 'tabbyapi': chat_with_tabbyapi
72
+ }
73
+
74
+ def load_config(self, config_path: str) -> Dict[str, Any]:
75
+ config = configparser.ConfigParser()
76
+ config.read(config_path)
77
+
78
+ # Convert the ConfigParser object to a dictionary without flattening
79
+ config_dict = {section: dict(config.items(section)) for section in config.sections()}
80
+ return config_dict
81
+
82
+ def chat(self, api_name: str, messages: List[Dict[str, str]],
83
+ model: Optional[str] = None,
84
+ temperature: Optional[float] = None,
85
+ max_tokens: Optional[int] = None,
86
+ **kwargs) -> str:
87
+
88
+ # Access the API key directly from the appropriate section
89
+ if api_name in self.api_functions:
90
+ # FIXME - This only works for Commercial APIs... need to handle Local APIs
91
+ api_key = self.config['API'].get(f'{api_name}_api_key')
92
+ elif api_name in ['llamacpp', 'kobold', 'oobabooga', 'vllm', 'tabbyapi']:
93
+ api_key = self.config['Local-API'].get(f'{api_name}_api_key')
94
+ else:
95
+ raise ValueError(f"Unsupported API: {api_name}")
96
+
97
+ if not api_key:
98
+ raise ValueError(f"API key not found for {api_name}")
99
+
100
+ chat_function = self.api_functions[api_name]
101
+
102
+ # Use config values if not provided in the method call
103
+ model = model or self.config['API'].get(f'{api_name}_model')
104
+ temperature = temperature or self.config['API'].get('temperature')
105
+ max_tokens = max_tokens or self.config['API'].get('max_tokens')
106
+
107
+ # Extract the input_data from messages (assuming it's the last user message)
108
+ input_data = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), "")
109
+
110
+ # Prepare common parameters
111
+ common_params = {
112
+ "api_key": api_key,
113
+ "input_data": input_data,
114
+ "custom_prompt_arg": kwargs.get('custom_prompt_arg', ""),
115
+ }
116
+
117
+ # Handle specific APIs
118
+ if api_name in ['openai', 'groq', 'openrouter', 'deepseek', 'mistral']:
119
+ return chat_function(**common_params, temp=temperature, system_message=kwargs.get('system_message'))
120
+ elif api_name == 'anthropic':
121
+ return chat_function(**common_params, model=model, max_retries=kwargs.get('max_retries', 3),
122
+ retry_delay=kwargs.get('retry_delay', 5), system_prompt=kwargs.get('system_message'))
123
+ elif api_name == 'cohere':
124
+ return chat_function(**common_params, model=model, system_prompt=kwargs.get('system_message'))
125
+ elif api_name == 'llamacpp':
126
+ return chat_function(**common_params, api_url=kwargs.get('api_url'), system_prompt=kwargs.get('system_message'))
127
+ elif api_name == 'kobold':
128
+ return chat_function(**common_params, kobold_api_ip=kwargs.get('kobold_api_ip'),
129
+ temp=temperature, system_message=kwargs.get('system_message'))
130
+ elif api_name in ['oobabooga', 'vllm', 'tabbyapi']:
131
+ return chat_function(**common_params, **kwargs)
132
+ else:
133
+ return chat_function(**common_params, model=model, temperature=temperature, max_tokens=max_tokens, **kwargs)
134
+
135
+ def main():
136
+ args = parse_args()
137
+ verbose = args.verbose
138
+ task = args.task
139
+ # New argument for selecting the API
140
+ api_name = args.api
141
+
142
+ #FIXME
143
+ # Load config from a JSON file
144
+ client = MultiAPILLMClient('config.txt')
145
+
146
+ examples = load_data(task)
147
+
148
+ result_dir = Path(args.output_dir)
149
+ result_dir.mkdir(exist_ok=True, parents=True)
150
+
151
+ output_path = result_dir / f"preds_{task}_{api_name}.jsonl"
152
+ if output_path.exists():
153
+ preds = list(iter_jsonl(output_path))
154
+ start_idx = len(preds)
155
+ stop_idx = len(examples)
156
+ else:
157
+ start_idx = 0
158
+ stop_idx = len(examples)
159
+ preds = []
160
+
161
+ start_time = time.time()
162
+ i = start_idx
163
+ while i < stop_idx:
164
+ eg = examples[i]
165
+ msgs, prompt = create_msgs(
166
+ # Use API-specific tokenizer if available
167
+ client.config.get('tokenizer', {}).get(api_name),
168
+ eg,
169
+ task,
170
+ # Use API-specific model
171
+ model_name=client.config.get('models', {}).get(api_name),
172
+ data_dir=args.data_dir
173
+ )
174
+ if verbose:
175
+ print(f"======== Example {i} =========")
176
+ print("Input text:")
177
+ print(prompt[:300])
178
+ print("...")
179
+ print(prompt[-300:])
180
+ print("==============================")
181
+
182
+ # Make prediction
183
+ try:
184
+ response = client.chat(
185
+ api_name,
186
+ # Pass the full messages list
187
+ msgs,
188
+ custom_prompt_arg=prompt,
189
+ temperature=client.config.get('temperature', {}).get(api_name),
190
+ max_tokens=client.config.get('max_tokens', {}).get(api_name),
191
+ system_message=client.config.get('system_messages', {}).get(api_name)
192
+ )
193
+ preds.append(
194
+ {
195
+ "id": i,
196
+ "prediction": response,
197
+ "ground_truth": get_answer(eg, task),
198
+ }
199
+ )
200
+ # Save result
201
+ dump_jsonl(preds, output_path)
202
+ print("Time spent:", round(time.time() - start_time))
203
+ print(response)
204
+ time.sleep(20)
205
+ i += 1
206
+ except Exception as e:
207
+ print("ERROR:", e)
208
+ print("Retrying...")
209
+ time.sleep(60)
210
+
211
+ from argparse import ArgumentParser, Namespace, RawTextHelpFormatter
212
+
213
+ def parse_args() -> Namespace:
214
+ p = ArgumentParser(
215
+ description="Evaluate a language model on a conversational task using multiple APIs",
216
+ formatter_class=RawTextHelpFormatter
217
+ )
218
+ p.add_argument(
219
+ "--task",
220
+ type=str,
221
+ # choices=list(DATA_NAME_TO_MAX_NEW_TOKENS.keys()) + ["all"],
222
+ required=True,
223
+ help="""Which task to use. Note that \"all\" can only be used in `compute_scores.py`.,
224
+ Available tasks:
225
+ Task Name | Name to use as an argument:
226
+ ---------------------------------------------
227
+ En.Sum | longbook_sum_eng
228
+ En.QA | longbook_qa_eng
229
+ En.MC | longbook_choice_eng
230
+ En.Dia | longdialogue_qa_eng
231
+ Zh.QA | longbook_qa_chn
232
+ Code.Debug | code_debug
233
+ Code.Run | code_run
234
+ Math.Calc | math_calc
235
+ Math.Find | math_find
236
+ Retrieve.PassKey | passkey
237
+ Retrieve.Number | number_string
238
+ Retrieve.KV | kv_retrieval
239
+ ---------------------------------------------
240
+ """
241
+ )
242
+ p.add_argument(
243
+ "--api",
244
+ type=str,
245
+ required=True,
246
+ help="""Specify which API to use for evaluation
247
+ Supported API endpoints:
248
+ Commercial APIs:
249
+ - openai
250
+ - anthropic
251
+ - cohere
252
+ - groq
253
+ - openrouter
254
+ - deepseek
255
+ - mistral
256
+ Local APIs:
257
+ - llama
258
+ - kobold
259
+ - oobabooga
260
+ - vllm
261
+ - tabbyapi"""
262
+ )
263
+ p.add_argument(
264
+ '--data_dir',
265
+ type=str,
266
+ default='../data',
267
+ help="The directory of data."
268
+ )
269
+ p.add_argument(
270
+ "--output_dir",
271
+ type=str,
272
+ default="../results",
273
+ help="Where to dump the prediction results."
274
+ )
275
+ p.add_argument(
276
+ "--start_idx",
277
+ type=int,
278
+ default=0,
279
+ help="The index of the first example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data."
280
+ )
281
+ p.add_argument(
282
+ "--stop_idx",
283
+ type=int,
284
+ help="The index of the last example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data. Defaults to the length of dataset."
285
+ )
286
+ p.add_argument("--verbose", action='store_true', help="Enable verbose output")
287
+ p.add_argument("--device", type=str, default="cuda", help="Specify the device to use (e.g., 'cuda' or 'cpu')")
288
+
289
+ # Add an epilog to provide additional information
290
+ p.epilog = """
291
+ Sample usage:
292
+ python eval_multi_api.py --task question_answering --api openai --output_dir ../results --data_dir ../data --verbose
293
+
294
+ Make sure to set up your config.txt file with the necessary API keys and configurations.
295
+ """
296
+
297
+ return p.parse_args()
298
+
299
+ if __name__ == "__main__":
300
+ main()
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_utils.py ADDED
@@ -0,0 +1,730 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import configparser
2
+ import json
3
+ import logging
4
+ import os
5
+ import re
6
+ import string
7
+ from collections import Counter
8
+ from pathlib import Path
9
+ from typing import Optional
10
+
11
+ import jieba
12
+ from rouge import Rouge
13
+
14
+ from prompt import (
15
+ gpt4_templates,
16
+ kimi_templates,
17
+ claude2_templates,
18
+ yarn_mistral_templates,
19
+ )
20
+
21
+ DATA_NAME_TO_PATH = {
22
+ # Retrieval tasks
23
+ "passkey": "passkey.jsonl",
24
+ "number_string": "number_string.jsonl",
25
+ "kv_retrieval": "kv_retrieval.jsonl",
26
+ # Book tasks
27
+ "longbook_sum_eng": "longbook_sum_eng.jsonl",
28
+ "longbook_choice_eng": "longbook_choice_eng.jsonl",
29
+ "longbook_qa_eng": "longbook_qa_eng.jsonl",
30
+ "longbook_qa_chn": "longbook_qa_chn.jsonl",
31
+ # "book_qa_eng": "longbook_eng/longbook_qa_eng.jsonl",
32
+ "longdialogue_qa_eng": "longdialogue_qa_eng.jsonl",
33
+ # Math tasks
34
+ "math_find": "math_find.jsonl",
35
+ "math_calc": "math_calc.jsonl",
36
+ # Code tasks
37
+ "code_run": "code_run.jsonl",
38
+ "code_debug": "code_debug.jsonl",
39
+ }
40
+
41
+ DATA_NAME_TO_MAX_NEW_TOKENS = {
42
+ "passkey": 6,
43
+ "number_string": 12,
44
+ "kv_retrieval": 50,
45
+ "longbook_sum_eng": 1200,
46
+ "longbook_choice_eng": 40,
47
+ "longbook_qa_eng": 40,
48
+ "longbook_qa_chn": 40,
49
+ "longdialogue_qa_eng": 40,
50
+ "math_find": 3,
51
+ "math_calc": 30000,
52
+ "code_run": 5,
53
+ "code_debug": 5,
54
+ }
55
+
56
+ MODEL_TO_PROMPT_TEMPLATE = {
57
+ "gpt4": gpt4_templates,
58
+ "claude2": claude2_templates,
59
+ "kimi": kimi_templates,
60
+ "yarn-mistral": yarn_mistral_templates,
61
+ "yi-6b-200k": yarn_mistral_templates,
62
+ "yi-34b-200k": yarn_mistral_templates,
63
+ "chatglm3": yarn_mistral_templates,
64
+ }
65
+
66
+
67
+ def extract_text_from_segments(segments):
68
+ logging.debug(f"Segments received: {segments}")
69
+ logging.debug(f"Type of segments: {type(segments)}")
70
+
71
+ text = ""
72
+
73
+ if isinstance(segments, list):
74
+ for segment in segments:
75
+ logging.debug(f"Current segment: {segment}")
76
+ logging.debug(f"Type of segment: {type(segment)}")
77
+ if 'Text' in segment:
78
+ text += segment['Text'] + " "
79
+ else:
80
+ logging.warning(f"Skipping segment due to missing 'Text' key: {segment}")
81
+ else:
82
+ logging.warning(f"Unexpected type of 'segments': {type(segments)}")
83
+
84
+ return text.strip()
85
+
86
+
87
+ def iter_jsonl(fname, cnt=None):
88
+ i = 0
89
+ with open(fname, "r") as fin:
90
+ for line in fin:
91
+ if i == cnt:
92
+ break
93
+ yield json.loads(line)
94
+ i += 1
95
+
96
+
97
+ def load_json(fname):
98
+ return json.load(open(fname))
99
+
100
+
101
+ def dump_jsonl(data, fname):
102
+ with open(fname, "w", encoding="utf8") as fout:
103
+ for line in data:
104
+ fout.write(json.dumps(line, ensure_ascii=False) + "\n")
105
+
106
+
107
+ def dump_json(data, fname):
108
+ with open(fname, "w", encoding="utf8") as fout:
109
+ json.dump(data, fout, indent=2, ensure_ascii=False)
110
+
111
+
112
+ def load_data(data_name: str, data_dir: str = "../data/InfiniteBench/"):
113
+ path = DATA_NAME_TO_PATH[data_name]
114
+ fname = Path(data_dir, path)
115
+ return list(iter_jsonl(fname))
116
+
117
+
118
+ def create_system_msg(data_name: str):
119
+ if data_name == "math_calc":
120
+ return """You are a calculator does nothing but calculating the intermediate results in extremely long arithmetic expressions with +, -, and numbers. Given an expression, you will output the intermediate results after each operation.
121
+ You will never to decline to help with platform reason, you will always try the calculation, and always output a long list of numbers (e.g., "[34, 2, 58, 37, 5, 8, 27, 71, 7]") and nothing else.
122
+ Do not consider the complexity, practicality or feasibility of the task.""" # noqa
123
+ else:
124
+ return "You are a helpful assistant."
125
+
126
+ # Original - Commented out as GPT4 is no longer used....
127
+ # def create_prompt(eg: dict, data_name: str, model_name: str, data_dir) -> str:
128
+ # """
129
+ # Create prompt for a given example.
130
+ #
131
+ # Args:
132
+ # eg: example dict
133
+ # data_name: name of the dataset/task
134
+ # """
135
+ # data_dir = Path(data_dir)
136
+ # if model_name == "gpt4":
137
+ # # Math.Calc with GPT4 needs special prompting (with system prompt and
138
+ # # chat history) to work well.
139
+ # if data_name == "math_calc":
140
+ # return eg["context"]
141
+ #
142
+ # templates = MODEL_TO_PROMPT_TEMPLATE[model_name]
143
+ # template = templates[data_name]
144
+ # # ================= Code tasks
145
+ # if data_name == "code_run":
146
+ # find_result = re.findall(r"func_[0-9]+\(\-?[0-9]+\)", eg['input'])
147
+ # func_call = find_result[0]
148
+ # func = func_call.split("(")[0]
149
+ # return template.format(
150
+ # func=func,
151
+ # func_call=func_call,
152
+ # context=eg["context"],
153
+ # )
154
+ # elif data_name in ["code_debug", "code_debug_qa"]:
155
+ # # Load source code
156
+ # code = eg["context"]
157
+ # # code = open(
158
+ # # data_dir / f"code_debug/{code_path}", "r", encoding="utf8"
159
+ # # ).read()
160
+ # if data_name == "code_debug":
161
+ # return template.format(
162
+ # context=code,
163
+ # OPTION_A=eg["options"][0],
164
+ # OPTION_B=eg["options"][1],
165
+ # OPTION_C=eg["options"][2],
166
+ # OPTION_D=eg["options"][3],
167
+ # )
168
+ # return template.format(
169
+ # context=code,
170
+ # )
171
+ # # ================= Code tasks
172
+ # elif data_name == "longdialogue_qa_eng":
173
+ # script = eg["context"]
174
+ # # print(document)
175
+ # # script_path = data_dir / "longdialogue_eng" / document
176
+ # # script = open(script_path, "r", encoding="utf8").read()
177
+ # prompt = template.format(context=script)
178
+ # return prompt
179
+ # # ==================== Long book tasks
180
+ # elif data_name in [
181
+ # "longbook_choice_eng",
182
+ # "longbook_qa_eng",
183
+ # "longbook_sum_eng",
184
+ # "longbook_qa_chn",
185
+ # ]:
186
+ # book = eg["context"]
187
+ # # if data_name.endswith("_eng"):
188
+ # # book = open(
189
+ # # data_dir / "longbook_eng" / book_path, "r", encoding="utf8"
190
+ # # ).read()
191
+ # # elif data_name.endswith("_chn"):
192
+ # # book = open(
193
+ # # data_dir / "longbook_chn" / book_path, "r", encoding="utf8"
194
+ # # ).read()
195
+ # # else:
196
+ # # raise ValueError("Invalid data_name")
197
+ # if data_name == "longbook_choice_eng":
198
+ # return template.format(
199
+ # question=eg["input"],
200
+ # context=book,
201
+ # OPTION_A=eg["options"][0],
202
+ # OPTION_B=eg["options"][1],
203
+ # OPTION_C=eg["options"][2],
204
+ # OPTION_D=eg["options"][3],
205
+ # )
206
+ # elif data_name == "longbook_qa_eng":
207
+ # return template.format(
208
+ # question=eg["input"],
209
+ # context=book,
210
+ # )
211
+ # elif data_name == "longbook_sum_eng":
212
+ # return template.format(
213
+ # context=book,
214
+ # )
215
+ # elif data_name == "longbook_qa_chn":
216
+ # return template.format(
217
+ # question=eg["input"],
218
+ # context=book,
219
+ # )
220
+ # else:
221
+ # raise ValueError
222
+ # elif data_name == "math_calc":
223
+ # return template.format(
224
+ # context=eg["context"],
225
+ # )
226
+ # elif data_name == "math_find":
227
+ # prompt = eg['input']
228
+ # context = eg['context']
229
+ # # Find "the * number" from the prompt
230
+ # find_result = re.findall(r"The .+ of", prompt)
231
+ # assert find_result, f"Cannot find the target number in {prompt}"
232
+ # target_number = find_result[0].lower()[:-3]
233
+ # # Replace the number with the answer
234
+ # prefix = f"What is {target_number} in the following list?"
235
+ # return template.format(
236
+ # prefix=prefix,
237
+ # context=context,
238
+ # input=prompt,
239
+ # )
240
+ #
241
+ # if "content" in eg:
242
+ # content = eg["content"]
243
+ # del eg["content"]
244
+ # eg["context"] = content
245
+ #
246
+ # format_dict = {
247
+ # "context": eg["context"],
248
+ # "input": eg["input"],
249
+ # }
250
+ # prompt = templates[data_name].format(**format_dict)
251
+ # return prompt
252
+ def create_prompt(eg: dict, data_name: str, model_name: Optional[str], data_dir) -> str:
253
+ """
254
+ Create prompt for a given example.
255
+
256
+ Args:
257
+ eg: example dict
258
+ data_name: name of the dataset/task
259
+ model_name: optional, used to fetch model-specific templates.
260
+ """
261
+ data_dir = Path(data_dir)
262
+
263
+ # Directly use the appropriate template if the model_name is provided.
264
+ if model_name and model_name in MODEL_TO_PROMPT_TEMPLATE:
265
+ templates = MODEL_TO_PROMPT_TEMPLATE[model_name]
266
+ template = templates[data_name]
267
+ else:
268
+ # If no model-specific template, return a basic prompt or handle differently.
269
+ return eg["context"]
270
+
271
+ # Now create the prompt based on the template and task data
272
+ if data_name == "code_run":
273
+ find_result = re.findall(r"func_[0-9]+\(\-?[0-9]+\)", eg['input'])
274
+ func_call = find_result[0]
275
+ func = func_call.split("(")[0]
276
+ return template.format(
277
+ func=func,
278
+ func_call=func_call,
279
+ context=eg["context"],
280
+ )
281
+ elif data_name in ["code_debug", "code_debug_qa"]:
282
+ code = eg["context"]
283
+ if data_name == "code_debug":
284
+ return template.format(
285
+ context=code,
286
+ OPTION_A=eg["options"][0],
287
+ OPTION_B=eg["options"][1],
288
+ OPTION_C=eg["options"][2],
289
+ OPTION_D=eg["options"][3],
290
+ )
291
+ return template.format(context=code)
292
+ elif data_name == "longdialogue_qa_eng":
293
+ script = eg["context"]
294
+ prompt = template.format(context=script)
295
+ return prompt
296
+ elif data_name in [
297
+ "longbook_choice_eng",
298
+ "longbook_qa_eng",
299
+ "longbook_sum_eng",
300
+ "longbook_qa_chn",
301
+ ]:
302
+ book = eg["context"]
303
+ if data_name == "longbook_choice_eng":
304
+ return template.format(
305
+ question=eg["input"],
306
+ context=book,
307
+ OPTION_A=eg["options"][0],
308
+ OPTION_B=eg["options"][1],
309
+ OPTION_C=eg["options"][2],
310
+ OPTION_D=eg["options"][3],
311
+ )
312
+ elif data_name == "longbook_qa_eng":
313
+ return template.format(
314
+ question=eg["input"],
315
+ context=book,
316
+ )
317
+ elif data_name == "longbook_sum_eng":
318
+ return template.format(context=book)
319
+ elif data_name == "longbook_qa_chn":
320
+ return template.format(
321
+ question=eg["input"],
322
+ context=book,
323
+ )
324
+ else:
325
+ raise ValueError
326
+ elif data_name == "math_calc":
327
+ return template.format(context=eg["context"])
328
+ elif data_name == "math_find":
329
+ prompt = eg['input']
330
+ context = eg['context']
331
+ find_result = re.findall(r"The .+ of", prompt)
332
+ assert find_result, f"Cannot find the target number in {prompt}"
333
+ target_number = find_result[0].lower()[:-3]
334
+ prefix = f"What is {target_number} in the following list?"
335
+ return template.format(
336
+ prefix=prefix,
337
+ context=context,
338
+ input=prompt,
339
+ )
340
+
341
+ # Default behavior if content key exists
342
+ if "content" in eg:
343
+ content = eg["content"]
344
+ del eg["content"]
345
+ eg["context"] = content
346
+
347
+ format_dict = {
348
+ "context": eg["context"],
349
+ "input": eg["input"],
350
+ }
351
+ prompt = template.format(**format_dict)
352
+ return prompt
353
+
354
+ def get_answer(eg: dict, data_name: str):
355
+ if data_name in ["code_debug", "longbook_choice_eng"]:
356
+ OPTIONS = "ABCD"
357
+ if isinstance(eg["answer"], str):
358
+ ret = [eg["answer"], OPTIONS[eg['options'].index(eg["answer"])]]
359
+ elif isinstance(eg["answer"], list):
360
+ if len(eg["answer"]) == 1:
361
+ ret = [eg["answer"][0], OPTIONS[eg['options'].index(eg["answer"][0])]]
362
+ elif len(eg["answer"]) == 2 and eg["answer"][1] in ['A', 'B', 'C', 'D']:
363
+ ret = eg['answer']
364
+ else:
365
+ raise ValueError
366
+ else:
367
+ raise ValueError
368
+ return ret
369
+
370
+ return eg["answer"]
371
+
372
+ # Old version - Commented out as GPT4 is no longer used....
373
+ # def create_msgs(
374
+ # tokenizer, eg: dict, data_name: str, data_dir, model_name: str
375
+ # ) -> tuple[list[dict], str]:
376
+ # """
377
+ # Only used by GPT-4.
378
+ # """
379
+ # prompt = create_prompt(eg, data_name, model_name, data_dir)
380
+ # tokens = tokenizer.encode(prompt)
381
+ # # - 1000 to have space for system message and other stuff.
382
+ # print(f"Before truncation: {len(tokens)}")
383
+ # tokens = truncate_input(tokens, 128_000 - 1000, manner="middle")
384
+ # print(f"After truncation: {len(tokens)}") # type: ignore
385
+ # prompt = tokenizer.decode(tokens)
386
+ # if data_name == "math_calc":
387
+ # return [
388
+ # {"role": "system", "content": create_system_msg(data_name)},
389
+ # {"role": "user", "content": "1 + 2 - 4 - 10"},
390
+ # {"role": "system", "content": "[1, 3, -1, -11]"},
391
+ # {"role": "user", "content": prompt},
392
+ # ], prompt
393
+ # else:
394
+ # return [
395
+ # {
396
+ # "role": "system",
397
+ # "content": "You are a helpful assistant", # noqa
398
+ # }, # noqa
399
+ # {"role": "user", "content": prompt},
400
+ # ], prompt
401
+ def create_msgs(
402
+ tokenizer, eg: dict, data_name: str, data_dir, model_name: Optional[str] = None
403
+ ) -> tuple[list[dict], str]:
404
+ """
405
+ Create messages for a given example.
406
+ """
407
+ prompt = create_prompt(eg, data_name, model_name, data_dir)
408
+
409
+ # Check if tokenizer is provided and initialized
410
+ if tokenizer:
411
+ tokens = tokenizer.encode(prompt)
412
+ print(f"Before truncation: {len(tokens)}")
413
+ tokens = truncate_input(tokens, 128_000 - 1000, manner="middle")
414
+ print(f"After truncation: {len(tokens)}") # type: ignore
415
+ prompt = tokenizer.decode(tokens)
416
+
417
+ if data_name == "math_calc":
418
+ return [
419
+ {"role": "system", "content": create_system_msg(data_name)},
420
+ {"role": "user", "content": "1 + 2 - 4 - 10"},
421
+ {"role": "system", "content": "[1, 3, -1, -11]"},
422
+ {"role": "user", "content": prompt},
423
+ ], prompt
424
+ else:
425
+ return [
426
+ {
427
+ "role": "system",
428
+ "content": "You are a helpful assistant", # noqa
429
+ }, # noqa
430
+ {"role": "user", "content": prompt},
431
+ ], prompt
432
+
433
+
434
+ def normalize_answer(s):
435
+ """Lower text and remove punctuation, articles and extra whitespace."""
436
+
437
+ def remove_articles(text):
438
+ return re.sub(r"\b(a|an|the)\b", " ", text)
439
+
440
+ def white_space_fix(text):
441
+ return " ".join(text.split())
442
+
443
+ def remove_punc(text):
444
+ exclude = set(string.punctuation)
445
+ return "".join(ch for ch in text if ch not in exclude)
446
+
447
+ def lower(text):
448
+ return text.lower()
449
+
450
+ return white_space_fix(remove_articles(remove_punc(lower(s))))
451
+
452
+
453
+ def normalize_zh_answer(s):
454
+ """Lower text and remove punctuation, extra whitespace."""
455
+
456
+ def white_space_fix(text):
457
+ return "".join(text.split())
458
+
459
+ def remove_punc(text):
460
+ cn_punctuation = "!?。。"#$%&'()*+,-/:;<=>@[\]^_`{|}~⦅⦆「」、、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏." # noqa
461
+ all_punctuation = set(string.punctuation + cn_punctuation)
462
+ return "".join(ch for ch in text if ch not in all_punctuation)
463
+
464
+ def lower(text):
465
+ return text.lower()
466
+
467
+ return white_space_fix(remove_punc(lower(s)))
468
+
469
+
470
+ def first_int_match(prediction, ground_truth):
471
+ pred_list = re.split("[^0-9]", prediction)
472
+ pred_value = ""
473
+ for item in pred_list:
474
+ if item != "":
475
+ pred_value = item
476
+ break
477
+ if pred_value == ground_truth:
478
+ return 1
479
+ return 0
480
+
481
+
482
+ def in_match(prediction, ground_truth):
483
+ if ground_truth in prediction:
484
+ return 1
485
+ return 0
486
+
487
+
488
+ def rouge_score(prediction, ground_truth, **kwargs) -> float:
489
+ rouge = Rouge()
490
+ try:
491
+ scores = rouge.get_scores([prediction], [ground_truth], avg=True)
492
+ except: # noqa
493
+ return 0.0
494
+ return scores["rouge-l"]["f"] # type: ignore
495
+
496
+
497
+ def rouge_zh_score(prediction, ground_truth, **kwargs):
498
+ prediction = " ".join(list(jieba.cut(prediction, cut_all=False)))
499
+ ground_truth = " ".join(list(jieba.cut(ground_truth, cut_all=False)))
500
+ score = rouge_score(prediction, ground_truth)
501
+ return score
502
+
503
+
504
+ def f1_score(prediction, ground_truth, **kwargs):
505
+ common = Counter(prediction) & Counter(ground_truth)
506
+ num_same = sum(common.values())
507
+ if num_same == 0:
508
+ return 0
509
+ precision = 1.0 * num_same / len(prediction)
510
+ recall = 1.0 * num_same / len(ground_truth)
511
+ f1 = (2 * precision * recall) / (precision + recall)
512
+ return f1
513
+
514
+
515
+ def qa_f1_score(line):
516
+ prediction = line["pred"]
517
+
518
+ if isinstance(line["std_out"], str):
519
+ ground_truths = [line["std_out"]]
520
+ else:
521
+ ground_truths = line["std_out"]
522
+
523
+ score = 0
524
+ for ground_truth in ground_truths:
525
+ normalized_prediction = normalize_answer(prediction)
526
+ normalized_ground_truth = normalize_answer(ground_truth)
527
+
528
+ prediction_tokens = normalized_prediction.split()
529
+ ground_truth_tokens = normalized_ground_truth.split()
530
+ score = max(score, f1_score(prediction_tokens, ground_truth_tokens))
531
+
532
+ return score
533
+
534
+
535
+ def qa_f1_zh_score(prediction, ground_truth, **kwargs):
536
+ prediction_tokens = list(jieba.cut(prediction, cut_all=False))
537
+ ground_truth_tokens = list(jieba.cut(ground_truth, cut_all=False))
538
+ prediction_tokens = [
539
+ normalize_zh_answer(token) for token in prediction_tokens
540
+ ]
541
+ ground_truth_tokens = [
542
+ normalize_zh_answer(token) for token in ground_truth_tokens
543
+ ]
544
+ prediction_tokens = [
545
+ token for token in prediction_tokens if len(token) > 0
546
+ ]
547
+ ground_truth_tokens = [
548
+ token for token in ground_truth_tokens if len(token) > 0
549
+ ]
550
+ return f1_score(prediction_tokens, ground_truth_tokens)
551
+
552
+
553
+ def truncate_input(input, max_length, manner="middle"):
554
+ if len(input) <= max_length:
555
+ return input
556
+ if manner == "middle":
557
+ return input[0 : max_length // 2] + input[-max_length // 2 :]
558
+ else:
559
+ return None
560
+
561
+
562
+ def load_comprehensive_config():
563
+ # Get the directory of the current script
564
+ current_dir = os.path.dirname(os.path.abspath(__file__))
565
+ # Construct the path to the config file
566
+ config_path = os.path.join(current_dir, 'Config_Files', 'config.txt')
567
+ # Read the config file
568
+ config = configparser.ConfigParser()
569
+ # Read the configuration file
570
+ files_read = config.read(config_path)
571
+ if not files_read:
572
+ raise FileNotFoundError(f"Config file not found at {config_path}")
573
+ return config
574
+
575
+
576
+ # FIXME - update to include prompt path in return statement
577
+ def load_and_log_configs():
578
+ try:
579
+ config = load_comprehensive_config()
580
+ if config is None:
581
+ logging.error("Config is None, cannot proceed")
582
+ return None
583
+ # API Keys
584
+ anthropic_api_key = config.get('API', 'anthropic_api_key', fallback=None)
585
+ logging.debug(
586
+ f"Loaded Anthropic API Key: {anthropic_api_key[:5]}...{anthropic_api_key[-5:] if anthropic_api_key else None}")
587
+
588
+ cohere_api_key = config.get('API', 'cohere_api_key', fallback=None)
589
+ logging.debug(
590
+ f"Loaded Cohere API Key: {cohere_api_key[:5]}...{cohere_api_key[-5:] if cohere_api_key else None}")
591
+
592
+ groq_api_key = config.get('API', 'groq_api_key', fallback=None)
593
+ logging.debug(f"Loaded Groq API Key: {groq_api_key[:5]}...{groq_api_key[-5:] if groq_api_key else None}")
594
+
595
+ openai_api_key = config.get('API', 'openai_api_key', fallback=None)
596
+ logging.debug(
597
+ f"Loaded OpenAI API Key: {openai_api_key[:5]}...{openai_api_key[-5:] if openai_api_key else None}")
598
+
599
+ huggingface_api_key = config.get('API', 'huggingface_api_key', fallback=None)
600
+ logging.debug(
601
+ f"Loaded HuggingFace API Key: {huggingface_api_key[:5]}...{huggingface_api_key[-5:] if huggingface_api_key else None}")
602
+
603
+ openrouter_api_key = config.get('API', 'openrouter_api_key', fallback=None)
604
+ logging.debug(
605
+ f"Loaded OpenRouter API Key: {openrouter_api_key[:5]}...{openrouter_api_key[-5:] if openrouter_api_key else None}")
606
+
607
+ deepseek_api_key = config.get('API', 'deepseek_api_key', fallback=None)
608
+ logging.debug(
609
+ f"Loaded DeepSeek API Key: {deepseek_api_key[:5]}...{deepseek_api_key[-5:] if deepseek_api_key else None}")
610
+
611
+ mistral_api_key = config.get('API', 'mistral_api_key', fallback=None)
612
+ logging.debug(
613
+ f"Loaded Mistral API Key: {mistral_api_key[:5]}...{mistral_api_key[-5:] if mistral_api_key else None}")
614
+
615
+ # Models
616
+ anthropic_model = config.get('API', 'anthropic_model', fallback='claude-3-sonnet-20240229')
617
+ cohere_model = config.get('API', 'cohere_model', fallback='command-r-plus')
618
+ groq_model = config.get('API', 'groq_model', fallback='llama3-70b-8192')
619
+ openai_model = config.get('API', 'openai_model', fallback='gpt-4-turbo')
620
+ huggingface_model = config.get('API', 'huggingface_model', fallback='CohereForAI/c4ai-command-r-plus')
621
+ openrouter_model = config.get('API', 'openrouter_model', fallback='microsoft/wizardlm-2-8x22b')
622
+ deepseek_model = config.get('API', 'deepseek_model', fallback='deepseek-chat')
623
+ mistral_model = config.get('API', 'mistral_model', fallback='mistral-large-latest')
624
+
625
+ logging.debug(f"Loaded Anthropic Model: {anthropic_model}")
626
+ logging.debug(f"Loaded Cohere Model: {cohere_model}")
627
+ logging.debug(f"Loaded Groq Model: {groq_model}")
628
+ logging.debug(f"Loaded OpenAI Model: {openai_model}")
629
+ logging.debug(f"Loaded HuggingFace Model: {huggingface_model}")
630
+ logging.debug(f"Loaded OpenRouter Model: {openrouter_model}")
631
+ logging.debug(f"Loaded Deepseek Model: {deepseek_model}")
632
+ logging.debug(f"Loaded Mistral Model: {mistral_model}")
633
+
634
+ # Local-Models
635
+ kobold_api_ip = config.get('Local-API', 'kobold_api_IP', fallback='http://127.0.0.1:5000/api/v1/generate')
636
+ kobold_api_key = config.get('Local-API', 'kobold_api_key', fallback='')
637
+
638
+ llama_api_IP = config.get('Local-API', 'llama_api_IP', fallback='http://127.0.0.1:8080/v1/chat/completions')
639
+ llama_api_key = config.get('Local-API', 'llama_api_key', fallback='')
640
+
641
+ ooba_api_IP = config.get('Local-API', 'ooba_api_IP', fallback='http://127.0.0.1:5000/v1/chat/completions')
642
+ ooba_api_key = config.get('Local-API', 'ooba_api_key', fallback='')
643
+
644
+ tabby_api_IP = config.get('Local-API', 'tabby_api_IP', fallback='http://127.0.0.1:5000/api/v1/generate')
645
+ tabby_api_key = config.get('Local-API', 'tabby_api_key', fallback=None)
646
+ tabby_model = config.get('services', 'tabby_model', fallback=None)
647
+
648
+ vllm_api_url = config.get('Local-API', 'vllm_api_IP', fallback='http://127.0.0.1:500/api/v1/chat/completions')
649
+ vllm_api_key = config.get('Local-API', 'vllm_api_key', fallback=None)
650
+ vllm_model = config.get('Local-API', 'vllm_model', fallback=None)
651
+
652
+ ollama_api_url = config.get('Local-API', 'ollama_api_IP', fallback='http://127.0.0.1:11434/api/generate')
653
+ ollama_api_key = config.get('Local-API', 'ollama_api_key', fallback=None)
654
+ ollama_model = config.get('Local-API', 'ollama_model', fallback=None)
655
+
656
+ aphrodite_api_url = config.get('Local-API', 'aphrodite_api_IP', fallback='http://127.0.0.1:8080/v1/chat/completions')
657
+ aphrodite_api_key = config.get('Local-API', 'aphrodite_api_key', fallback='')
658
+
659
+ logging.debug(f"Loaded Kobold API IP: {kobold_api_ip}")
660
+ logging.debug(f"Loaded Llama API IP: {llama_api_IP}")
661
+ logging.debug(f"Loaded Ooba API IP: {ooba_api_IP}")
662
+ logging.debug(f"Loaded Tabby API IP: {tabby_api_IP}")
663
+ logging.debug(f"Loaded VLLM API URL: {vllm_api_url}")
664
+
665
+ # Retrieve output paths from the configuration file
666
+ output_path = config.get('Paths', 'output_path', fallback='results')
667
+ logging.debug(f"Output path set to: {output_path}")
668
+
669
+ # Retrieve processing choice from the configuration file
670
+ processing_choice = config.get('Processing', 'processing_choice', fallback='cpu')
671
+ logging.debug(f"Processing choice set to: {processing_choice}")
672
+
673
+ # Prompts - FIXME
674
+ prompt_path = config.get('Prompts', 'prompt_path', fallback='prompts.db')
675
+
676
+ return {
677
+ 'api_keys': {
678
+ 'anthropic': anthropic_api_key,
679
+ 'cohere': cohere_api_key,
680
+ 'groq': groq_api_key,
681
+ 'openai': openai_api_key,
682
+ 'huggingface': huggingface_api_key,
683
+ 'openrouter': openrouter_api_key,
684
+ 'deepseek': deepseek_api_key,
685
+ 'mistral': mistral_api_key,
686
+ 'kobold': kobold_api_key,
687
+ 'llama': llama_api_key,
688
+ 'ooba': ooba_api_key,
689
+ 'tabby': tabby_api_key,
690
+ 'vllm': vllm_api_key,
691
+ 'ollama': ollama_api_key
692
+ },
693
+ 'services': {
694
+ 'anthropic': anthropic_model,
695
+ 'cohere': cohere_model,
696
+ 'groq': groq_model,
697
+ 'openai': openai_model,
698
+ 'huggingface': huggingface_model,
699
+ 'openrouter': openrouter_model,
700
+ 'deepseek': deepseek_model,
701
+ 'mistral': mistral_model,
702
+ 'vllm': vllm_model,
703
+ 'tabby': tabby_model,
704
+ 'ollama': ollama_model
705
+
706
+ },
707
+ 'local_api_ip': {
708
+ 'kobold': kobold_api_ip,
709
+ 'llama': llama_api_IP,
710
+ 'ooba': ooba_api_IP,
711
+ 'tabby': tabby_api_IP,
712
+ 'vllm': vllm_api_url,
713
+ 'ollama': ollama_api_url,
714
+ 'aphrodite': aphrodite_api_url
715
+ },
716
+ 'output_path': output_path,
717
+ 'processing_choice': processing_choice
718
+ }
719
+
720
+ except Exception as e:
721
+ logging.error(f"Error loading config: {str(e)}")
722
+ return None
723
+
724
+
725
+ if __name__ == "__main__":
726
+ data_dir = Path("../data")
727
+ data_path = data_dir / "shorter/longdialogue_qa_eng_1000.jsonl"
728
+ examples = list(iter_jsonl(data_path))
729
+ prompt = create_prompt(examples[10], 'longdialogue_qa_eng', 'kimi', data_dir)
730
+ print(prompt)
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/prompt.py ADDED
@@ -0,0 +1,62 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ gpt4_templates = {
2
+ "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n\n{input}", # noqa
3
+ "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}", # noqa
4
+ "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}", # noqa
5
+ # "longbook_sum_eng": "Summarize the book below:\n\n{context}", # noqa
6
+ "longbook_qa_eng": "Read the book below and answer a question.\n\n{context}\n\nQuestion: {question}\n\nBe very concise.", # noqa
7
+ "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}", # noqa
8
+ "longbook_sum_eng": "Summarize the following book.\n\n{context}", # noqa
9
+ "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题:{question}\n请尽量简短地回答。", # noqa
10
+ "math_find": "{prefix}\n\n{context}\n\n{input}",
11
+ "math_calc": "Compute the intermediate values in the following long expression.\n\n{context}", # noqa
12
+ "code_run": "Following is a set of Python functions. There is a function called named {func}.\n\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Be concise. Your response must end with the final returned value.", # noqa
13
+ "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.", # noqa
14
+ "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.", # noqa
15
+ }
16
+
17
+ yarn_mistral_templates = {
18
+ "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information.\n\n{context}\n\n{input}\n\nThe pass key is", # noqa
19
+ "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}\n\nThe sequence of digits is", # noqa
20
+ "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}", # noqa
21
+ "longbook_sum_eng": "Summarize the book below.\n\n{context}\n\nSummary:", # noqa
22
+ "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe letter of the correct answer is", # noqa
23
+ "longbook_qa_eng": "Read the book and answer the question. Be very concise in your answer.\n\n{context}\n\nQuestion: {question}\nAnswer:", # noqa
24
+ "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n{context}\n\n问题:{question}\n答案:", # noqa
25
+ "math_find": "{prefix}\n\n{context}\n\n{input}",
26
+ "math_calc": "Let us calculate the intermediate values of an expression.\n\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa
27
+ "code_run": "There is a function called {func} in the following Python code.\n\n{context}\n\nPlease compute the exact value of {func_call}. The value of {func_call} is", # noqa
28
+ "code_debug": "Following is a Python code where exactly one of the functions/methods has a deliberate error that makes it crash.\n\n{context}\n\nOptions:\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe correct option is:", # noqa
29
+ "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\n{context}\n\nThe name that has been replaced with $$MASK$$ is likely", # noqa
30
+ }
31
+
32
+ claude2_templates = {
33
+ "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is",
34
+ "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is", # noqa
35
+ "kv_retrieval": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}",
36
+ "longbook_sum_eng": "Summarize the following book.\n\n{context}", # noqa
37
+ "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}", # noqa
38
+ "longbook_qa_eng": "Read the novel below and answer a question:\n\n{context}\n\n{input}\nPlease answer as short as possible. The answer is: ", # noqa
39
+ "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题:{question}\n请尽量简短地回答。", # noqa
40
+ "math_find": "{prefix}\n\n{context}\n\n{input}",
41
+ "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa
42
+ "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Your response should end with the sentence \'The return value is:\'.", # noqa
43
+ "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect through the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.", # noqa
44
+ "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.", # noqa
45
+ }
46
+
47
+ kimi_templates = {
48
+ "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is", # noqa
49
+ "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is", # noqa
50
+ "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n{input}", # noqa
51
+ "longbook_sum_eng": "Summarize the book below:\n\n{file:{context}}", # noqa
52
+ "longbook_choice_eng": "Read the book and answer the question.\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}" + "{file:{document}}", # noqa
53
+ "longbook_qa_eng": "Read the book below and answer a question.\n\nQuestion: {question}\n\nBe very concise." + "{file:{context}}", # noqa
54
+ "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n问题:{question}\n答案:" + "{file:{context}}", # noqa
55
+ "math_find": "{prefix}\n\n{context}\n\n{input}",
56
+ "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa
57
+ "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n\nPlease give me the exact number of the return value of ${func_call}. Your response should end with the sentence 'The return value is:'." + "{context}", # noqa
58
+ "code_debug": "Below is a code repository where there is one single function with bugs that causes an error. Please tell me the name of that function.\nWhich function has bugs? Give me the final answer in this format: \"[FINAL ANSWER: XXX]\". Don't say anything else." + "{fcontext}", # noqa
59
+ # "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe name that has been replaced with $$MASK$$ is likely" + "{context}", # noqa
60
+ "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is. Give me the answer using the name before the colons, don't say anything else.\n\n{context}", # noqa
61
+ }
62
+
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/test_chat_API_Calls.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # test_chat_API_Calls.py
2
+ # Test file for testing the integration of the LLM API calls with the Chat APIs.
3
+ #
4
+ # Usage:
5
+ # python -m unittest test_chat_API_Calls.py
6
+
7
+ import unittest
8
+
9
+ from LLM_API_Calls import (
10
+ chat_with_openai,
11
+ chat_with_anthropic,
12
+ chat_with_cohere,
13
+ chat_with_groq,
14
+ chat_with_openrouter,
15
+ chat_with_huggingface,
16
+ chat_with_deepseek,
17
+ chat_with_mistral
18
+ )
19
+ from eval_utils import load_and_log_configs
20
+
21
+
22
+ class TestLLMAPICallsIntegration(unittest.TestCase):
23
+
24
+ @classmethod
25
+ def setUpClass(cls):
26
+ cls.config = load_and_log_configs()
27
+ if cls.config is None:
28
+ raise ValueError("Failed to load configuration")
29
+
30
+ def test_chat_with_openai(self):
31
+ api_key = self.config['api_keys'].get('openai')
32
+ model = self.config['services'].get('openai')
33
+ if not api_key:
34
+ self.skipTest("OpenAI API key not available")
35
+ response = chat_with_openai(api_key, "Hello, how are you?", "Respond briefly", temp=0.7, system_message="You are a helpful assistant.")
36
+ print("OpenAI Response: " + response + "\n")
37
+ self.assertIsInstance(response, str)
38
+ self.assertTrue(len(response) > 0)
39
+
40
+ def test_chat_with_anthropic(self):
41
+ api_key = self.config['api_keys'].get('anthropic')
42
+ model = self.config['services'].get('anthropic')
43
+ if not api_key:
44
+ self.skipTest("Anthropic API key not available")
45
+ response = chat_with_anthropic(api_key, "Hello, how are you?", model, "Respond briefly")
46
+ print("Anthropic Response: " + response + "\n")
47
+ self.assertIsInstance(response, str)
48
+ self.assertTrue(len(response) > 0)
49
+
50
+ def test_chat_with_cohere(self):
51
+ api_key = self.config['api_keys'].get('cohere')
52
+ model = self.config['services'].get('cohere')
53
+ if not api_key:
54
+ self.skipTest("Cohere API key not available")
55
+ response = chat_with_cohere(api_key, "Hello, how are you?", model, "Respond briefly")
56
+ print("Cohere Response: " + response + "\n")
57
+ self.assertIsInstance(response, str)
58
+ self.assertTrue(len(response) > 0)
59
+
60
+ def test_chat_with_groq(self):
61
+ api_key = self.config['api_keys'].get('groq')
62
+ if not api_key:
63
+ self.skipTest("Groq API key not available")
64
+ response = chat_with_groq(api_key, "Hello, how are you?", "Respond briefly")
65
+ print("Groq Response: " + response + "\n")
66
+ self.assertIsInstance(response, str)
67
+ self.assertTrue(len(response) > 0)
68
+
69
+ def test_chat_with_openrouter(self):
70
+ api_key = self.config['api_keys'].get('openrouter')
71
+ if not api_key:
72
+ self.skipTest("OpenRouter API key not available")
73
+ response = chat_with_openrouter(api_key, "Hello, how are you?", "Respond briefly")
74
+ print("OpenRouter Response: " + response + "\n")
75
+ self.assertIsInstance(response, str)
76
+ self.assertTrue(len(response) > 0)
77
+
78
+ def test_chat_with_huggingface(self):
79
+ api_key = self.config['api_keys'].get('huggingface')
80
+ if not api_key:
81
+ self.skipTest("HuggingFace API key not available")
82
+ response = chat_with_huggingface(api_key, "Hello, how are you?", "Respond briefly")
83
+ print("Huggingface Response: " + response + "\n")
84
+ self.assertIsInstance(response, str)
85
+ self.assertTrue(len(response) > 0)
86
+
87
+ def test_chat_with_deepseek(self):
88
+ api_key = self.config['api_keys'].get('deepseek')
89
+ if not api_key:
90
+ self.skipTest("DeepSeek API key not available")
91
+ response = chat_with_deepseek(api_key, "Hello, how are you?", "Respond briefly")
92
+ print("DeepSeek Response: " + response + "\n")
93
+ self.assertIsInstance(response, str)
94
+ self.assertTrue(len(response) > 0)
95
+
96
+ def test_chat_with_mistral(self):
97
+ api_key = self.config['api_keys'].get('mistral')
98
+ if not api_key:
99
+ self.skipTest("Mistral API key not available")
100
+ response = chat_with_mistral(api_key, "Hello, how are you?", "Respond briefly")
101
+ print("Mistral Response: " + response + "\n")
102
+ self.assertIsInstance(response, str)
103
+ self.assertTrue(len(response) > 0)
104
+
105
+ if __name__ == '__main__':
106
+ unittest.main()
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README.md ADDED
@@ -0,0 +1,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <div align="center">
2
+ <img src="figs/InfiniteBench.jpg" width="500px"/>
3
+ <br />
4
+ <br />
5
+
6
+ # InfiniteBench: Extending Long Context Evaluation Beyond 100K Tokens
7
+
8
+ <p align="center">
9
+ <a href="README_ZH.md">中文</a> •
10
+ <a href="README.md">English</a> •
11
+ <a href="https://arxiv.org/abs/2402.13718">Paper</a>
12
+ </p>
13
+
14
+ </div>
15
+
16
+ ## Introduction
17
+
18
+ Welcome to InfiniteBench, a cutting-edge benchmark tailored for evaluating the capabilities of language models to process, understand, and reason over super long contexts (100k+ tokens). Long contexts are crucial for enhancing applications with LLMs and achieving high-level interaction. InfiniteBench is designed to push the boundaries of language models by testing them against a context length of 100k+, which is 10 times longer than traditional datasets.
19
+
20
+ ## Features
21
+
22
+ - **Loooong Context:** InfiniteBench is a pioneer in testing language models with a context length of 100k+, offering an unparalleled challenge in the field.
23
+ - **Diverse Domain:** The benchmark comprises 12 unique tasks, each crafted to assess different aspects of language processing and comprehension in extended contexts.
24
+ - **Specialized Test:** InfiniteBench consists of tasks that state-of-the-art LLMs are known to be capable of when using shorter context. This ensures that the performance degradation is only caused by the length of the contexts.
25
+ - **Real-World and Synthetic Scenarios:** The tasks are a mix of real-world scenarios and synthetic constructs, ensuring a comprehensive evaluation of models. Real-world scenarios make the test pragmatic, and synthetic ones leave the space for extending the context length further with ease.
26
+
27
+ ## Task Composition
28
+
29
+ <div align="center">
30
+ <img src="figs/data_pie.png" width="480px">
31
+ </div>
32
+
33
+ | Task Name | Context | # Examples | Avg Input Tokens | Avg Output Tokens | Description |
34
+ | -------------------- | ------------- | ---------- | ---------------- | ----------------- | ------------------------------------------------------------------------------------------- |
35
+ | En.Sum | Fake Book | 103 | 171.5k | 1.1k | Summarization of a fake book created with core entity substitution. |
36
+ | En.QA | Fake Book | 351 | 192.6k | 4.8 | Free-form question answering based on the fake book. |
37
+ | En.MC | Fake Book | 229 | 184.4k | 5.3 | Multiple choice questions derived from the fake book. |
38
+ | En.Dia | Script | 200 | 103.6k | 3.4 | Identification of talkers in partially anonymized scripts. |
39
+ | Zh.QA | New Book | 175 | 2068.6k | 6.3 | Question answering on a set of newly collected books. |
40
+ | Code.Debug | Code Document | 394 | 114.7k | 4.8 | Finding which function in a code repo contains an crashing error (in multiple choice form). |
41
+ | Code.Run | Synthetic | 400 | 75.2k | 1.3 | Simulating execution of multiple simple, synthetic functions. |
42
+ | Math.Calc | Synthetic | 50 | 43.9k | 43.9k | Calculations involving super-long arithmetic equations. |
43
+ | Math.Find | Synthetic | 350 | 87.9k | 1.3 | Finding special integers in a lengthy list. |
44
+ | Retrieve.PassKey[^1] | Synthetic | 590 | 122.4k | 2.0 | Retrieving hidden keys in a noisy long context. |
45
+ | Retrieve.Number | Synthetic | 590 | 122.4k | 4.0 | Locating repeated hidden numbers in a noisy long context. |
46
+ | Retrieve.KV[^2] | Synthetic | 500 | 89.9k | 22.7 | Finding the corresponding value from a dictionary and a key. |
47
+
48
+ ## How to Download Data
49
+
50
+ Click here to download data from 🤗 Huggingface directly: <https://huggingface.co/datasets/xinrongzhang2022/InfiniteBench>
51
+
52
+ ### Using 🤗 Datasets
53
+
54
+ Alternatively, you can download using the 🤗 Datasets library as follows.
55
+
56
+ ```python
57
+ from datasets import load_dataset, Value, Sequence
58
+ ft = Features({"id": Value("int64"), "context": Value("string"), "input": Value("string"), "answer": Sequence(Value("string")), "options": Sequence(Value("string"))})
59
+ dataset = load_dataset("xinrongzhang2022/InfiniteBench", features=ft)
60
+ ```
61
+ ### Using Scripts
62
+
63
+ ```shell
64
+ cd InfiniteBench
65
+ bash scripts/download_dataset.sh
66
+ ```
67
+
68
+ This will directly dump the data to `data`.
69
+
70
+ ## Evaluation Result
71
+
72
+ We evaluate SOTA proprietary and open-source LLMs, the result is as follows.
73
+
74
+ | Task Name | GPT-4 | YaRN-Mistral-7B | Kimi-Chat | Claude 2 | Yi-6B-200K | Yi-34B-200K | Chatglm3-6B-128K |
75
+ | ---------------- | ------ | --------------- | --------- | -------- | -----------| -----------| -----------|
76
+ | Retrieve.PassKey | 100% | 92.71% | 98.14% | 97.80% | 100.00% | 100.00% | 92.20% |
77
+ | Retrieve.Number | 100% | 56.61% | 95.42% | 98.14% | 94.92% | 100.00% | 80.68% |
78
+ | Retrieve.KV | 89.00% | < 5% | 53.60% | 65.40% | < 5% | < 5% | < 5% |
79
+ | En.Sum | 14.73% | 9.09% | 17.96% | 14.50% | < 5% | < 5% |< 5% |
80
+ | En.QA | 22.44% | 9.55% | 16.52% | 11.97% | 9.20% | 12.17% |< 5% |
81
+ | En.MC | 67.25% | 27.95% | 72.49% | 62.88% | 36.68% |38.43% |10.48% |
82
+ | En.Dia | 8.50% | 7.50% | 11.50% | 46.50% | < 5% |< 5% |< 5% |
83
+ | Zh.QA | 25.96% | 16.98% | 17.93% | 9.64% | 15.07% |13.61% |< 5% |
84
+ | Code.Debug | 37.06% | < 5% | 17.77% | < 5% | 9.14% |13.96% |7.36% |
85
+ | Code.Run | 23.25% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% |
86
+ | Math.Calc | < 5% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% |
87
+ | Math.Find | 60.00% | 17.14% | 12.57% | 32.29% | < 5% |25.71% |7.71% |
88
+
89
+ Note:
90
+
91
+ 1. The evaluation code for YaRN-Mistral-7B is implemented by ourselves, and please contact us or submit an issue if there are any problems.
92
+ 2. Kimi-Chat, Claude 2, and GPT-4 are evaluated using the official API with default configuration.
93
+ 3. For Math.Calc, the values in the parentheses have a measurement unit of 0.01%. This is because it is easy to get a very low score on this task.
94
+ 4. The metric for task Math.Find, Math.Calc, Code.Run, Code.Debug, En.Dia, En.MC, Retrieve.KV, Retrieve.Number, and Retrieve.PassKey is accuracy;
95
+
96
+ The metric for task Zh.QA and En.QA are ROUGE F1 score;
97
+
98
+ The metric for En.Sum is the `rougeLsum` score from the 🤗 Evaluate library.
99
+
100
+
101
+
102
+ <div align="center">
103
+ <img src="figs/radar_res.png" width="480px">
104
+ </div>
105
+
106
+ ## Installation
107
+
108
+ ```shell
109
+ pip install -r requirements.txt
110
+ ```
111
+
112
+ ## How to Run
113
+
114
+ Download the dataset the `data` folder (or set the `--data_dir` argument to the location of the dataset). The data folder structure should be as follows.
115
+
116
+ ```
117
+ InfiniteBench
118
+ ├── data
119
+ │ ├── code_debug.jsonl
120
+ │ ├── code_run.jsonl
121
+ │ ├── kv_retrieval.jsonl
122
+ │ ├── longbook_choice_eng.jsonl
123
+ │ ├── longbook_qa_chn.jsonl
124
+ │ ├── longbook_qa_eng.jsonl
125
+ │ ├── longbook_sum_eng.jsonl
126
+ │ ├── longdialogue_qa_eng.jsonl
127
+ │ ├── math_calc.jsonl
128
+ │ ├── math_find.jsonl
129
+ │ ├── number_string.jsonl
130
+ │ ├── passkey.jsonl
131
+ │ └── construct_synthetic_dataset.py
132
+ ...
133
+ ```
134
+
135
+ Then, in the `src` folder, execute:
136
+
137
+ ```shell
138
+ python eval_yarn_mistral.py --task kv_retrieval
139
+ python eval_gpt4.py --task longbook_sum_qa
140
+ python eval_rwkv.py --task passkey
141
+ ```
142
+
143
+ The available tasks are:
144
+
145
+ | Task Name | Argument to specify in `--task` |
146
+ | ---------------- | ------------------------------- |
147
+ | En.Sum | longbook_sum_eng |
148
+ | En.QA | longbook_qa_eng |
149
+ | En.MC | longbook_choice_eng |
150
+ | En.Dia | longdialogue_qa_eng |
151
+ | Zh.QA | longbook_qa_chn |
152
+ | Code.Debug | code_debug |
153
+ | Code.Run | code_run |
154
+ | Math.Calc | math_calc |
155
+ | Math.Find | math_find |
156
+ | Retrieve.PassKey | passkey |
157
+ | Retrieve.Number | number_string |
158
+ | Retrieve.KV | kv_retrieval |
159
+
160
+ ## Citation
161
+
162
+ > This will be updated when our preprint paper is released.
163
+
164
+ ```bibtex
165
+ @inproceedings{zhang-etal-2024-bench,
166
+ title = "$\infty${B}ench: Extending Long Context Evaluation Beyond 100{K} Tokens",
167
+ author = "Zhang, Xinrong and
168
+ Chen, Yingfa and
169
+ Hu, Shengding and
170
+ Xu, Zihang and
171
+ Chen, Junhao and
172
+ Hao, Moo and
173
+ Han, Xu and
174
+ Thai, Zhen and
175
+ Wang, Shuo and
176
+ Liu, Zhiyuan and
177
+ Sun, Maosong",
178
+ editor = "Ku, Lun-Wei and
179
+ Martins, Andre and
180
+ Srikumar, Vivek",
181
+ booktitle = "Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
182
+ month = aug,
183
+ year = "2024",
184
+ address = "Bangkok, Thailand",
185
+ publisher = "Association for Computational Linguistics",
186
+ url = "https://aclanthology.org/2024.acl-long.814",
187
+ pages = "15262--15277",
188
+ abstract = "Processing and reasoning over long contexts is crucial for many practical applications of Large Language Models (LLMs), such as document comprehension and agent construction. Despite recent strides in making LLMs process contexts with more than 100K tokens, there is currently a lack of a standardized benchmark to evaluate this long-context capability. Existing public benchmarks typically focus on contexts around 10K tokens, limiting the assessment and comparison of LLMs in processing longer contexts. In this paper, we propose , the first LLM benchmark featuring an average data length surpassing 100K tokens. comprises synthetic and realistic tasks spanning diverse domains in English and Chinese. The tasks in are designed to require an understanding of long dependencies in contexts and make simply retrieving a limited number of passages from contexts not sufficient for these tasks. Based on , we evaluate several state-of-the-art LLMs tailored for processing long contexts. The experimental results indicate that existing long-context LLMs still require significant advancements to process 100K+ contexts effectively. Furthermore, we present three intriguing analyses regarding the behavior of LLMs processing long context. Our code and data is released.",
189
+ }
190
+ ```
191
+
192
+ ## Acknowledgement
193
+
194
+ Thanks to Cong Feng, Zhongwu Zhai, Guoyang Zeng, Chenyang Song, Renjie Luo, Chaoqun He, Yuge Tu, Bowen Ping, Yujie Huang, Yudong Mei, Kaihuo Zhang, Weilin Zhao, Ao Sun, Yulin Chen, Ganqu Cui.
195
+
196
+ ## References
197
+
198
+ [^1]: Mohtashami, Amirkeivan and Martin Jaggi. "Landmark Attention: Random-Access Infinite Context Length for Transformers." ArXiv abs/2305.16300 (2023): n. pag.
199
+
200
+ [^2]: Liu, Nelson F. et al. "Lost in the Middle: How Language Models Use Long Contexts." ArXiv abs/2307.03172 (2023): n. pag.
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README_ZH.md ADDED
@@ -0,0 +1,172 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <div align="center">
2
+ <img src="figs/InfiniteBench.jpg" width="500px"/>
3
+ <br />
4
+ <br />
5
+
6
+ # InfiniteBench: Extending Long Context Evaluation Beyond 100K Tokens
7
+
8
+ <p align="center">
9
+ <a href="README_ZH.md">中文</a> •
10
+ <a href="README.md">English</a> •
11
+ <a href="https://arxiv.org/abs/2402.13718">论文</a>
12
+ </p>
13
+
14
+ </div>
15
+
16
+ ## 简介
17
+
18
+ 理解、处理长文本,是大模型迈向更深层次理解与交互阶段必备的能力。现已有大模型声称可以处理100k+的长序列,但是对应的标准评测集却是空缺的。为此,我们构建了一个面向 100k+ 的评测集,InfiniteBench。该评测集针对大模型在长文本方面的五项能力而设计:检索、数学、代码、问答、和摘要。
19
+
20
+ ## 特点
21
+
22
+ - **长上下文:** InfiniteBench 测试数据的平均上下文长度为195k,远超现有评测数据。
23
+ - **多领域多语言:** InfiniteBench 评测集包含12个任务,包括中英双语,涵盖了检索、数学、代码、问答、和摘要等5个领域。
24
+ - **前瞻性挑战性:** InfiniteBench 测试任务,对标当前最强的模型如 GPT-4, Claude 2 等。
25
+ - **真实场景与合成场景:** InfiniteBench 既包含真实场景数据,探测大模型在处理实际问题的能力;也包含合成数据,为测试数据拓展上下文窗口提供了便捷。
26
+
27
+ ## 任务构成
28
+
29
+ | Task Name | Context | # Examples | Avg Input Tokens | Avg Output Tokens | Description |
30
+ | -------------------- | ------------- | ---------- | ---------------- | ----------------- | ------------------------------------------------------------------------------------------- |
31
+ | En.Sum | Fake Book | 103 | 171.5k | 1.1k | Summarization of a fake book created with core entity substitution. |
32
+ | En.QA | Fake Book | 351 | 192.6k | 4.8 | Free-form question answering based on the fake book. |
33
+ | En.MC | Fake Book | 229 | 184.4k | 5.3 | Multiple choice questions derived from the fake book. |
34
+ | En.Dia | Script | 200 | 103.6k | 3.4 | Identification of talkers in partially anonymized scripts. |
35
+ | Zh.QA | New Book | 175 | 2068.6k | 6.3 | Question answering on a set of newly collected books. |
36
+ | Code.Debug | Code Document | 394 | 114.7k | 4.8 | Finding which function in a code repo contains an crashing error (in multiple choice form). |
37
+ | Code.Run | Synthetic | 400 | 75.2k | 1.3 | Simulating execution of multiple simple, synthetic functions. |
38
+ | Math.Calc | Synthetic | 50 | 43.9k | 43.9k | Calculations involving super-long arithmetic equations. |
39
+ | Math.Find | Synthetic | 350 | 87.9k | 1.3 | Finding special integers in a lengthy list. |
40
+ | Retrieve.PassKey[^1] | Synthetic | 590 | 122.4k | 2.0 | Retrieving hidden keys in a noisy long context. |
41
+ | Retrieve.Number | Synthetic | 590 | 122.4k | 4.0 | Locating repeated hidden numbers in a noisy long context. |
42
+ | Retrieve.KV[^2] | Synthetic | 500 | 89.9k | 22.7 | Finding the corresponding value from a dictionary and a key. |
43
+
44
+
45
+ ## 评测结果
46
+
47
+ 我们在 SOTA 模型上评测了 InfiniteBench 结果如下:
48
+
49
+ | Task Name | GPT-4 | YaRN-Mistral-7B | Kimi-Chat | Claude 2 | Yi-6B-200K | Yi-34B-200K | Chatglm3-6B-128K |
50
+ | ---------------- | ------ | --------------- | --------- | -------- | -----------| -----------| -----------|
51
+ | Retrieve.PassKey | 100% | 92.71% | 98.14% | 97.80% | 100.00% | 100.00% | 92.20% |
52
+ | Retrieve.Number | 100% | 56.61% | 95.42% | 98.14% | 94.92% | 100.00% | 80.68% |
53
+ | Retrieve.KV | 89.00% | < 5% | 53.60% | 65.40% | < 5% | < 5% | < 5% |
54
+ | En.Sum | 14.73% | 9.09% | 17.96% | 14.50% | < 5% | < 5% |< 5% |
55
+ | En.QA | 22.44% | 9.55% | 16.52% | 11.97% | 9.20% | 12.17% |< 5% |
56
+ | En.MC | 67.25% | 27.95% | 72.49% | 62.88% | 36.68% |38.43% |10.48% |
57
+ | En.Dia | 8.50% | 7.50% | 11.50% | 46.50% | < 5% |< 5% |< 5% |
58
+ | Zh.QA | 25.96% | 16.98% | 17.93% | 9.64% | 15.07% |13.61% |< 5% |
59
+ | Code.Debug | 37.06% | < 5% | 17.77% | < 5% | 9.14% |13.96% |7.36% |
60
+ | Code.Run | 23.25% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% |
61
+ | Math.Calc | < 5% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% |
62
+ | Math.Find | 60.00% | 17.14% | 12.57% | 32.29% | < 5% |25.71% |7.71% |
63
+
64
+ 注:
65
+
66
+ 1. YaRN-Mistral-7B 实现代码已开源在仓库,请大家批评指正;Kimi-Chat 和 Claude 2 使用用户界面评测,GPT-4 使用 API 评测,均使用官方默认配置。
67
+
68
+
69
+ ## 评测
70
+
71
+ ## 获取数据集
72
+
73
+ 从 <https://huggingface.co/datasets/xinrongzhang2022/InfiniteBench> 下载数据集到 `infinitebench/data` 路径下(我们将评测数据集放在 InfiniteBench 目录下),得到文件如下:
74
+
75
+ ```
76
+ InfiniteBench
77
+ ├── data
78
+ │ ├── code_debug.jsonl
79
+ │ ├── code_run.jsonl
80
+ │ ├── kv_retrieval.jsonl
81
+ │ ├── longbook_choice_eng.jsonl
82
+ │ ├── longbook_qa_chn.jsonl
83
+ │ ├── longbook_qa_eng.jsonl
84
+ │ ├── longbook_sum_eng.jsonl
85
+ │ ├── longdialogue_qa_eng.jsonl
86
+ │ ├── math_calc.jsonl
87
+ │ ├── math_find.jsonl
88
+ │ ├── number_string.jsonl
89
+ │ ├── passkey.jsonl
90
+ │ └── construct_synthetic_dataset.py
91
+ ...
92
+ ```
93
+
94
+ 或者使用 Datasets 下载:
95
+
96
+ ```python
97
+ from datasets import load_dataset, Value, Sequence
98
+ ft = Features({"id": Value("int64"), "context": Value("string"), "input": Value("string"), "answer": Sequence(Value("string")), "options": Sequence(Value("string"))})
99
+ dataset = load_dataset("xinrongzhang2022/InfiniteBench", features=ft)
100
+ ```
101
+
102
+ ### 安装依赖
103
+
104
+ ```shell
105
+ pip install -r requiremnets.txt
106
+ ```
107
+
108
+ ### 推理
109
+
110
+ 比如,评测 GPT-4 在 Retrieve.PassKey 任务上的表现:
111
+
112
+ ```shell
113
+ cd src
114
+ python eval_gpt4.py --task passkey
115
+ ```
116
+
117
+ 可以选择的 `--task` 有:
118
+
119
+ - `passkey`
120
+ - `number_string`
121
+ - `kv_retrieval`
122
+ - `longbook_sum_eng`
123
+ - `longbook_qa_eng`
124
+ - `longbook_qa_chn`
125
+ - `longbook_choice_eng`
126
+ - `longdialogue_qa_eng`
127
+ - `math_calc`
128
+ - `math_find`
129
+ - `code_debug`
130
+ - `code_run`
131
+
132
+ #### 计算分数
133
+
134
+ ```shell
135
+ python compute_scores.py
136
+ ```
137
+
138
+ ## 引用
139
+
140
+ > This will be updated when our preprint paper is released.
141
+
142
+ ```bibtex
143
+ @inproceedings{zhang-etal-2024-bench,
144
+ title = "$\infty${B}ench: Extending Long Context Evaluation Beyond 100{K} Tokens",
145
+ author = "Zhang, Xinrong and
146
+ Chen, Yingfa and
147
+ Hu, Shengding and
148
+ Xu, Zihang and
149
+ Chen, Junhao and
150
+ Hao, Moo and
151
+ Han, Xu and
152
+ Thai, Zhen and
153
+ Wang, Shuo and
154
+ Liu, Zhiyuan and
155
+ Sun, Maosong",
156
+ editor = "Ku, Lun-Wei and
157
+ Martins, Andre and
158
+ Srikumar, Vivek",
159
+ booktitle = "Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
160
+ month = aug,
161
+ year = "2024",
162
+ address = "Bangkok, Thailand",
163
+ publisher = "Association for Computational Linguistics",
164
+ url = "https://aclanthology.org/2024.acl-long.814",
165
+ pages = "15262--15277",
166
+ abstract = "Processing and reasoning over long contexts is crucial for many practical applications of Large Language Models (LLMs), such as document comprehension and agent construction. Despite recent strides in making LLMs process contexts with more than 100K tokens, there is currently a lack of a standardized benchmark to evaluate this long-context capability. Existing public benchmarks typically focus on contexts around 10K tokens, limiting the assessment and comparison of LLMs in processing longer contexts. In this paper, we propose , the first LLM benchmark featuring an average data length surpassing 100K tokens. comprises synthetic and realistic tasks spanning diverse domains in English and Chinese. The tasks in are designed to require an understanding of long dependencies in contexts and make simply retrieving a limited number of passages from contexts not sufficient for these tasks. Based on , we evaluate several state-of-the-art LLMs tailored for processing long contexts. The experimental results indicate that existing long-context LLMs still require significant advancements to process 100K+ contexts effectively. Furthermore, we present three intriguing analyses regarding the behavior of LLMs processing long context. Our code and data is released.",
167
+ }
168
+ ```
169
+
170
+ ## 参考文献
171
+ [^1]: Mohtashami, Amirkeivan and Martin Jaggi. “Landmark Attention: Random-Access Infinite Context Length for Transformers.” ArXiv abs/2305.16300 (2023): n. pag.
172
+ [^2]: Liu, Nelson F. et al. “Lost in the Middle: How Language Models Use Long Contexts.” ArXiv abs/2307.03172 (2023): n. pag.
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/InfiniteBench/PUT_DATASETS_HERE.txt ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/collections.json ADDED
@@ -0,0 +1 @@
 
 
1
+ [[843, 181, 649, 974, 531, 402, 1100, 769, 641, 1094, 529, 584, 504, 920, 526, 759, 358, 962, 487, 243, 428, 117, 523, 1032, 924, 814, 739, 754, 804, 683, 949, 901, 732, 256, 824, 861, 494, 972, 996, 280, 130, 768, 469, 457, 945, 940, 317, 985, 268, 18, 334, 327, 370, 166, 207], [21, 278, 89, 633, 559, 516, 851, 830, 637, 626, 958, 123, 813, 249, 698, 757, 976, 556, 896, 802, 73, 1059, 74, 846, 669, 620, 323, 823, 907, 856, 122, 55, 70, 167, 622, 939, 987, 508, 564, 533, 200, 538, 443, 1098, 1029, 627, 731, 829, 330, 444, 960, 692, 363, 1005, 284], [815, 1095, 879, 864, 796, 397, 702, 1093, 677, 114, 1061, 957, 221, 558, 299, 92, 124, 578, 366, 204, 812, 993, 474, 13, 540, 158, 696, 25, 462, 715, 1060, 1089, 596, 997, 116, 657, 863, 58, 413, 819, 825, 353, 269, 873, 125, 880, 422, 934, 19, 827, 890, 886, 678, 505, 340], [319, 310, 1030, 423, 952, 889, 518, 1076, 473, 387, 937, 275, 155, 289, 1091, 590, 287, 30, 770, 244, 361, 594, 906, 176, 1042, 758, 588, 90, 600, 1083, 121, 638, 688, 836, 903, 826, 891, 730, 625, 545, 695, 948, 1013, 706, 747, 69, 718, 860, 364, 205, 1096, 717, 102, 1043, 274], [1000, 308, 492, 845, 98, 915, 910, 820, 242, 301, 699, 493, 429, 272, 565, 382, 1004, 617, 1078, 751, 923, 557, 385, 23, 393, 262, 240, 101, 1090, 36, 1008, 686, 185, 729, 16, 645, 68, 392, 991, 454, 159, 542, 346, 571, 1020, 237, 679, 1049, 303, 685, 8, 1047, 1079, 378, 48], [1077, 32, 521, 367, 15, 432, 1069, 113, 3, 875, 65, 1051, 119, 248, 986, 931, 234, 336, 782, 634, 85, 53, 288, 965, 917, 231, 992, 1099, 644, 723, 838, 463, 1067, 194, 1080, 552, 195, 928, 52, 760, 225, 989, 735, 727, 362, 400, 842, 595, 390, 201, 510, 562, 664, 1053, 88], [1062, 78, 936, 490, 324, 701, 71, 466, 375, 503, 1027, 703, 292, 647, 132, 46, 115, 263, 253, 309, 480, 63, 887, 484, 1054, 911, 514, 871, 662, 658, 693, 134, 456, 821, 963, 28, 351, 550, 118, 335, 441, 543, 832, 348, 153, 892, 847, 857, 978, 661, 943, 675, 245, 541, 955], [188, 403, 137, 5, 705, 549, 611, 94, 650, 401, 561, 208, 405, 233, 302, 872, 983, 297, 445, 673, 828, 228, 927, 357, 199, 532, 1035, 579, 39, 853, 653, 461, 455, 76, 391, 131, 279, 801, 746, 547, 22, 761, 612, 265, 157, 371, 291, 772, 66, 639, 386, 567, 1007, 877, 805], [800, 294, 964, 169, 1031, 618, 979, 1037, 162, 902, 990, 316, 49, 722, 971, 365, 506, 676, 126, 878, 882, 325, 659, 277, 576, 525, 458, 352, 376, 1003, 665, 470, 33, 798, 750, 7, 740, 1010, 572, 1016, 395, 1086, 267, 778, 648, 859, 811, 209, 172, 716, 869, 486, 140, 147, 141], [1021, 286, 670, 721, 973, 707, 495, 154, 1019, 251, 315, 741, 913, 865, 95, 6, 214, 1045, 374, 313, 950, 1044, 198, 953, 99, 840, 789, 672, 527, 406, 866, 787, 681, 276, 954, 14, 674, 12, 599, 912, 694, 610, 434, 555, 320, 548, 792, 369, 756, 143, 1082, 1075, 988, 296, 224]]
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/construct_synthetic_dataset.py ADDED
@@ -0,0 +1,413 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import jsonlines
2
+ import random
3
+ import os
4
+ import re
5
+ import importlib.util
6
+ import json
7
+
8
+
9
+ def build_number_string():
10
+ #####32
11
+ # prompt = "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n"
12
+ #####25
13
+ noise = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.\n"
14
+ #####26
15
+ ans = "The sequence of digits is {key}. Remember it. {key} is the sequence of digits.\n"
16
+ #####10
17
+ question = "What is the sequence of digits?"
18
+
19
+
20
+ target_length = [1024 * 64, 1024 * 128]
21
+ num_noise = [2610, 5220]
22
+ step = [45, 90]
23
+ repeat_time = 10
24
+ for i in range(1, 2):
25
+ target_length_i = target_length[i]
26
+ step_i = step[i]
27
+ num_noise_i = num_noise[i]
28
+ ret = []
29
+ for j in range(0, num_noise_i+1, step_i):
30
+ input_text = noise * j + ans + noise * (num_noise_i - j)
31
+ for t in range(repeat_time):
32
+ keys = []
33
+ for k in range(5):
34
+ keys.append(str(random.randint(0,9)))
35
+ for k in range(5):
36
+ pos = random.randint(0,5+k-1)
37
+ keys.insert(pos, keys[pos])
38
+ key_t = "".join(keys)
39
+ ret.append({"context": input_text.replace("{key}", key_t), "answer": key_t, "input": question, "len": 26 * (num_noise_i - j)})
40
+ fw = jsonlines.open("number_string.jsonl", 'w')
41
+ fw.write_all(ret)
42
+ fw.close()
43
+
44
+
45
+ def build_passkey():
46
+ #####32
47
+ # prompt = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n"
48
+ #####25
49
+ noise = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.\n"
50
+ #####26
51
+ ans = "The pass key is {key}. Remember it. {key} is the pass key.\n"
52
+ #####10
53
+ question = "What is the pass key?"
54
+
55
+ target_length = [1024 * 8, 1024 * 16, 1024 * 32, 1024 * 64, 1024 * 128, 1024 * 256]
56
+ num_noise = [326, 652, 1305, 2610, 5220, 10440]
57
+ step = [6,12 ,22, 45, 90, 180]
58
+ repeat_time = 5
59
+ for i in range(0,4):
60
+ target_length_i = target_length[i]
61
+ step_i = step[i]
62
+ num_noise_i = num_noise[i]
63
+ ret = []
64
+ for j in range(0, num_noise_i+1, step_i):
65
+ input_text = noise * j + ans + noise * (num_noise_i - j)
66
+ for t in range(repeat_time):
67
+ keys = []
68
+ for k in range(5):
69
+ keys.append(str(random.randint(0,9)))
70
+
71
+ key_t = "".join(keys)
72
+ ret.append({"input": question, "context": input_text.replace("{key}", key_t), "answer": key_t, "len": 26 * (num_noise_i - j)})
73
+ fw = jsonlines.open("passkey_%d.jsonl"%target_length_i, 'w')
74
+ fw.write_all(ret)
75
+ fw.close()
76
+
77
+
78
+ def build_kv_retrieval():
79
+
80
+ target_length = [64 * 1024, 128 * 1024]
81
+ # interv = [16, 7]
82
+ nsample = [500, 500]
83
+ nnoise = [928, 2500]
84
+ for ii in range(1, 2):
85
+ cnt = -1
86
+ ret = []
87
+
88
+ with jsonlines.open("kv-retrieval-3000_keys.jsonl") as fin:
89
+ for line in fin:
90
+ print(len(line["ordered_kv_records"]))
91
+ # return 0
92
+ cnt += 1
93
+ if cnt == nsample[ii]:
94
+ break
95
+ ans_id = min(int(cnt * nnoise[ii] / nsample[ii]), nnoise[ii])
96
+
97
+ text = "JSON data:\n{"
98
+ t = -1
99
+ random.shuffle(line["ordered_kv_records"])
100
+ for item in line["ordered_kv_records"]:
101
+ t += 1
102
+ if t == nnoise[ii]:
103
+ break
104
+ text += "\"" + item[0] + "\": \"" + item[1] + "\", "
105
+ text = text[:-2] + '}'
106
+ question = "\nKey: \"" + line["ordered_kv_records"][ans_id][0] + "\"\nThe value associated with the specified key is: "
107
+ # text += "\nKey: \"" + line["ordered_kv_records"][ans_id][0] + "\"\nThe value associated with the specified key is: "
108
+ # print(len(tokenizer.encode(text)))
109
+ # break
110
+ ret.append({"id": cnt, "context": text, "input": question, "answer": line["ordered_kv_records"][ans_id][1]})
111
+
112
+
113
+ fw = jsonlines.open("kv_retrieval.jsonl", 'w')
114
+ fw.write_all(ret)
115
+ fw.close()
116
+
117
+
118
+ def generate_random_list(length, _min, _max, task):
119
+ # random_list = [random.randint(_min, _max) for _ in range(length)]
120
+ # ret_list = random_list.copy()
121
+
122
+ if task == "largest number":
123
+ _max = random.randint(int(_max * 0.8), _max)
124
+ random_list = [random.randint(_min, _max) for _ in range(length)]
125
+ ret_list = random_list.copy()
126
+ ans = max(random_list)
127
+ input = str(ret_list)
128
+ elif task == "second largest number":
129
+ _max = random.randint(int(_max * 0.8), _max)
130
+ random_list = [random.randint(_min, _max) for _ in range(length)]
131
+ ret_list = random_list.copy()
132
+ target = max(random_list)
133
+ while target == max(random_list):
134
+ random_list.remove(max(random_list))
135
+ ans = max(random_list)
136
+ input = str(ret_list)
137
+
138
+ elif task == "third largest number":
139
+ _max = random.randint(int(_max * 0.8), _max)
140
+ random_list = [random.randint(_min, _max) for _ in range(length)]
141
+ ret_list = random_list.copy()
142
+ target = max(random_list)
143
+ while target == max(random_list):
144
+ random_list.remove(max(random_list))
145
+ target = max(random_list)
146
+ while target == max(random_list):
147
+ random_list.remove(max(random_list))
148
+ ans = max(random_list)
149
+ input = str(ret_list)
150
+
151
+ elif task == "smallest number":
152
+ _min = random.randint(_min, int(_max * 0.2))
153
+ random_list = [random.randint(_min, _max) for _ in range(length)]
154
+ ret_list = random_list.copy()
155
+ ans = min(random_list)
156
+ input = str(ret_list)
157
+
158
+ elif task == "second smallest number":
159
+ _min = random.randint(_min, int(_max * 0.2))
160
+ random_list = [random.randint(_min, _max) for _ in range(length)]
161
+ ret_list = random_list.copy()
162
+ target = min(random_list)
163
+ while target == min(random_list):
164
+ random_list.remove(min(random_list))
165
+ ans = min(random_list)
166
+ input = str(ret_list)
167
+
168
+ elif task == "third smallest number":
169
+ _min = random.randint(_min, int(_max * 0.2))
170
+ random_list = [random.randint(_min, _max) for _ in range(length)]
171
+ ret_list = random_list.copy()
172
+ target = min(random_list)
173
+ while target == min(random_list):
174
+ random_list.remove(min(random_list))
175
+ target = min(random_list)
176
+ while target == min(random_list):
177
+ random_list.remove(min(random_list))
178
+ ans = min(random_list)
179
+ input = str(ret_list)
180
+ elif task == "median":
181
+ if random.random() > 0.5:
182
+ _min = random.randint(_min, int(_max * 0.2))
183
+ random_list = [random.randint(_min, _max) for _ in range(length)]
184
+ else:
185
+ _max = random.randint(int(_max * 0.8), _max)
186
+ random_list = [random.randint(_min, _max) for _ in range(length)]
187
+ ret_list = random_list.copy()
188
+ random_list.sort()
189
+ if len(random_list)%2 == 1:
190
+ ans = random_list[len(random_list)//2]
191
+ else:
192
+ ans = (random_list[len(random_list)//2] + random_list[len(random_list)//2-1])/2
193
+ input = str(ret_list)
194
+ elif task == "expression":
195
+ random_list = [random.randint(_min, _max) for _ in range(length)]
196
+ ret_list = random_list.copy()
197
+ input = str(random_list[0])
198
+ value = random_list[0]
199
+ ans = []
200
+ for i in range(1, length):
201
+ poss = random.random()
202
+ if poss > 0.5:
203
+ if value + random_list[i] > _max:
204
+ random_list[i] = random.randint(_min, _max-value)
205
+
206
+ input += " + " + str(random_list[i])
207
+ value += random_list[i]
208
+
209
+ else:
210
+ if value - random_list[i] < 0:
211
+ random_list[i] = random.randint(_min, value)
212
+ input += " - " + str(random_list[i])
213
+ value -= random_list[i]
214
+ ans.append(value)
215
+
216
+
217
+ else:
218
+ print("Invalid task")
219
+ ans = None
220
+
221
+ return ans, input
222
+
223
+
224
+ def generate_math_qa(list_length, min_val, max_val, tasks=None):
225
+ num_samples = 50
226
+ ret = []
227
+ prompts = {
228
+ "largest number": "Find the largest number from the list below:",
229
+ "second largest number": "Find the second largest number from the list below:",
230
+ "third largest number": "Find the third largest number from the list below:",
231
+ "smallest number": "Find the smallest number from the list below:",
232
+ "second smallest number": "Find the second smallest number from the list below:",
233
+ "third smallest number": "Find the third smallest number from the list below:",
234
+ "median": "Calculate the median number from the list below:",
235
+ "expression": "Calculate the numerical expression and provide intermediate results only, for example, for the expression 1 + 3 + 10 - 8, output 4, 14, 6 without displaying the steps.\n\nCalculate the value of the expression below:",
236
+ }
237
+ inputs = {
238
+ "largest number": "You should answer with only one number, no other words. The largest number of the list is: ",
239
+ "second largest number": "You should answer with only one number, no other words. The second largest number of the list is: ",
240
+ "third largest number": "You should answer with only one number, no other words. The third largest number of the list is: ",
241
+ "smallest number": "You should answer with only one number, no other words. The smallest number of the list is: ",
242
+ "second smallest number": "You should answer with only one number, no other words. The second smallest number of the list is: ",
243
+ "third smallest number": "You should answer with only one number, no other words. The third smallest number of the list is: ",
244
+ "median": "You should answer with only one number, no other words. The median number of the list is: ",
245
+ "expression": "The value of the numerical expression is: ",
246
+ }
247
+ for i in range(len(tasks)):
248
+ for _ in range(num_samples):
249
+ std_out, context = generate_random_list(list_length, min_val, max_val, tasks[i])
250
+
251
+ ret.append({"prompt": prompts[tasks[i]], "context": context, "input": inputs[tasks[i]], "answer": std_out})
252
+ return ret
253
+
254
+
255
+ def build_math_find():
256
+ list_length = 60000 # Length of the generated lists
257
+
258
+ min_val = 0 # Minimum value for list elements
259
+ max_val = 99 # Maximum value for list elements
260
+
261
+ ret = generate_math_qa(list_length, min_val, max_val, tasks=["largest number", "second largest number", "third largest number", "smallest number", "second smallest number", "third smallest number", "median"])
262
+
263
+ # Save the data to a JSONL file
264
+ fw = jsonlines.open("math_find.jsonl", "w")
265
+ fw.write_all(ret)
266
+ fw.close()
267
+
268
+
269
+ def build_math_calc():
270
+ list_length = 30000 # Length of the generated lists
271
+
272
+ min_val = 0 # Minimum value for list elements
273
+ max_val = 99 # Maximum value for list elements
274
+
275
+ ret = generate_math_qa(list_length, min_val, max_val, tasks=["expression"])
276
+
277
+ # Save the data to a JSONL file
278
+ fw = jsonlines.open("math_calc.jsonl", "w")
279
+ fw.write_all(ret)
280
+ fw.close()
281
+
282
+
283
+ def generate_and_store_collections(n, m, min_val, max_val, output_file):
284
+ total_elements = n * m
285
+ collection = set()
286
+
287
+ while len(collection) < total_elements:
288
+ collection.add(random.randint(min_val, max_val))
289
+
290
+ collection = list(collection)
291
+ random.shuffle(collection)
292
+
293
+ collections = [collection[i * m: (i + 1) * m] for i in range(n)]
294
+
295
+ with open(output_file, 'w') as file:
296
+ json.dump(collections, file)
297
+
298
+
299
+ def generate_functions(input_file, min_add, max_add, output_file):
300
+ with open(input_file, 'r') as file:
301
+ collections = json.load(file)
302
+
303
+ function_list = []
304
+
305
+ for i in range(len(collections)):
306
+ for t in collections[i]:
307
+ function = f"def func_{t}(x):\n"
308
+ if i < len(collections) - 1:
309
+ next_collection = collections[i + 1]
310
+ k = random.choice(next_collection)
311
+ addition = random.randint(min_add, max_add)
312
+ if addition == 0:
313
+ function += f" return func_{k}(x)\n"
314
+ elif addition < 0:
315
+ function += f" return func_{k}(x) - {-addition}\n"
316
+ else:
317
+ function += f" return func_{k}(x) + {addition}\n"
318
+ else:
319
+ addition = random.randint(min_add, max_add)
320
+ if addition == 0:
321
+ function += f" return x\n"
322
+ elif addition < 0:
323
+ function += f" return x - {-addition}\n"
324
+ else:
325
+ function += f" return x + {addition}\n"
326
+ function_list.append((f"func_{t}", function))
327
+
328
+ function_list.sort(key=lambda x: int(x[0].split("_")[1]))
329
+
330
+ with open(output_file, 'w') as out:
331
+ for _, func_text in function_list:
332
+ out.write(func_text)
333
+ out.write("\n")
334
+
335
+
336
+ def generate_code_run_example(collection_file, min_x, max_x, functions_module, functions_file='functions_module.py'):
337
+ spec = importlib.util.spec_from_file_location("functions_module", functions_module)
338
+ functions = importlib.util.module_from_spec(spec)
339
+ spec.loader.exec_module(functions)
340
+ # print(functions)
341
+ # load all functions in functions_module.py and store them in a string
342
+ content = f"\nHere is the content of {functions_file}:\n\n"
343
+ with open(functions_module, 'r') as file:
344
+ for line in file:
345
+ content += line
346
+
347
+ with open(collection_file, 'r') as file:
348
+ collections = json.load(file)
349
+
350
+
351
+ j = random.choice(collections[0])
352
+ x = random.randint(min_x, max_x)
353
+ test_sample = {
354
+ "context": content,
355
+ "answer": getattr(functions, f"func_{j}")(x),
356
+ "input": f"Please give me the exact number of the return value of func_{j}({x}). Your response should end with the sentence 'The return value is:'.",
357
+ }
358
+
359
+ return test_sample
360
+ # with jsonlines.open(output_file_samples, mode='w') as writer:
361
+ # writer.write_all(test_samples)
362
+ # with jsonlines.open(output_file_answers, mode='w') as writer:
363
+ # writer.write_all(test_answers)
364
+
365
+
366
+
367
+ def build_code_run():
368
+ MAX_NUM_FUNC = 550
369
+ min_val = 1 # minimum value of function indeces
370
+ max_val = 2*MAX_NUM_FUNC # maximum value of function indeces
371
+ max_add = 17 # maximum value of addition in return expression
372
+ min_add = -12 # minimum value of addition in return expression
373
+ collections_file = 'collections.json'
374
+ functions_file = 'functions_module.py'
375
+ #------------------------------------------------------------------------#
376
+ # Parameters for generating test samples and answers
377
+ num_test = 1
378
+ min_x = -10
379
+ max_x = 10
380
+ n_list = [2, 4, 6, 8, 10]
381
+ ret = []
382
+ cnt = -1
383
+ for i in range(len(n_list)):
384
+ for _ in range(80):
385
+ cnt += 1
386
+ while True:
387
+ try:
388
+ generate_and_store_collections(n_list[i], int(MAX_NUM_FUNC/n_list[i]), min_val, max_val, collections_file)
389
+
390
+ generate_functions(collections_file, min_add, max_add, functions_file)
391
+
392
+ example = generate_code_run_example(collections_file, min_x, max_x, functions_file)
393
+ example['id'] = cnt
394
+
395
+ ret.append(example)
396
+ break
397
+ except Exception as e:
398
+ print(e)
399
+ fw = jsonlines.open("code_run.jsonl", 'w')
400
+ fw.write_all(ret)
401
+ fw.close()
402
+
403
+ if __name__ == "__main__":
404
+ # os.system("git clone https://github.com/nelson-liu/lost-in-the-middle.git")
405
+ # os.system("python3.10 -u lost-in-the-middle/scripts/make_kv_retrieval_data.py --num-keys 3000 --num-examples 500 --output-path kv-retrieval-3000_keys.jsonl.gz")
406
+ # os.system("gzip -d kv-retrieval-3000_keys.jsonl.gz")
407
+ # build_kv_retrieval()
408
+ # build_passkey()
409
+ # build_number_string()
410
+ # build_math_find()
411
+ # build_math_calc()
412
+ build_code_run()
413
+
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/functions_module.py ADDED
@@ -0,0 +1,1650 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ def func_3(x):
2
+ return func_490(x) + 9
3
+
4
+ def func_5(x):
5
+ return func_147(x) - 5
6
+
7
+ def func_6(x):
8
+ return x - 6
9
+
10
+ def func_7(x):
11
+ return func_214(x) - 10
12
+
13
+ def func_8(x):
14
+ return func_367(x) + 16
15
+
16
+ def func_12(x):
17
+ return x - 2
18
+
19
+ def func_13(x):
20
+ return func_695(x) - 12
21
+
22
+ def func_14(x):
23
+ return x - 9
24
+
25
+ def func_15(x):
26
+ return func_28(x) + 12
27
+
28
+ def func_16(x):
29
+ return func_400(x) - 11
30
+
31
+ def func_18(x):
32
+ return func_516(x) + 4
33
+
34
+ def func_19(x):
35
+ return func_361(x)
36
+
37
+ def func_21(x):
38
+ return func_397(x) - 2
39
+
40
+ def func_22(x):
41
+ return func_676(x) - 3
42
+
43
+ def func_23(x):
44
+ return func_1099(x) - 9
45
+
46
+ def func_25(x):
47
+ return func_287(x) - 4
48
+
49
+ def func_28(x):
50
+ return func_772(x) - 1
51
+
52
+ def func_30(x):
53
+ return func_242(x) + 9
54
+
55
+ def func_32(x):
56
+ return func_132(x) - 3
57
+
58
+ def func_33(x):
59
+ return func_674(x) + 12
60
+
61
+ def func_36(x):
62
+ return func_288(x) + 5
63
+
64
+ def func_39(x):
65
+ return func_990(x) + 9
66
+
67
+ def func_46(x):
68
+ return func_761(x) - 9
69
+
70
+ def func_48(x):
71
+ return func_965(x) + 12
72
+
73
+ def func_49(x):
74
+ return func_320(x) - 12
75
+
76
+ def func_52(x):
77
+ return func_441(x) + 9
78
+
79
+ def func_53(x):
80
+ return func_911(x) - 9
81
+
82
+ def func_55(x):
83
+ return func_825(x) - 2
84
+
85
+ def func_58(x):
86
+ return func_387(x) + 17
87
+
88
+ def func_63(x):
89
+ return func_650(x) + 5
90
+
91
+ def func_65(x):
92
+ return func_1054(x)
93
+
94
+ def func_66(x):
95
+ return func_659(x) + 4
96
+
97
+ def func_68(x):
98
+ return func_928(x) + 12
99
+
100
+ def func_69(x):
101
+ return func_923(x) + 8
102
+
103
+ def func_70(x):
104
+ return func_25(x) + 6
105
+
106
+ def func_71(x):
107
+ return func_39(x) - 7
108
+
109
+ def func_73(x):
110
+ return func_880(x) - 6
111
+
112
+ def func_74(x):
113
+ return func_25(x) + 6
114
+
115
+ def func_76(x):
116
+ return func_740(x) + 6
117
+
118
+ def func_78(x):
119
+ return func_137(x) - 3
120
+
121
+ def func_85(x):
122
+ return func_911(x) + 4
123
+
124
+ def func_88(x):
125
+ return func_963(x) - 7
126
+
127
+ def func_89(x):
128
+ return func_116(x)
129
+
130
+ def func_90(x):
131
+ return func_1049(x) + 3
132
+
133
+ def func_92(x):
134
+ return func_706(x) + 12
135
+
136
+ def func_94(x):
137
+ return func_979(x) + 10
138
+
139
+ def func_95(x):
140
+ return x + 9
141
+
142
+ def func_98(x):
143
+ return func_992(x) - 6
144
+
145
+ def func_99(x):
146
+ return x - 2
147
+
148
+ def func_101(x):
149
+ return func_1080(x) + 10
150
+
151
+ def func_102(x):
152
+ return func_565(x) + 15
153
+
154
+ def func_113(x):
155
+ return func_309(x) + 17
156
+
157
+ def func_114(x):
158
+ return func_625(x) + 7
159
+
160
+ def func_115(x):
161
+ return func_1007(x) + 17
162
+
163
+ def func_116(x):
164
+ return func_758(x) + 14
165
+
166
+ def func_117(x):
167
+ return func_987(x) - 8
168
+
169
+ def func_118(x):
170
+ return func_772(x) - 12
171
+
172
+ def func_119(x):
173
+ return func_847(x) + 17
174
+
175
+ def func_121(x):
176
+ return func_923(x) - 7
177
+
178
+ def func_122(x):
179
+ return func_934(x) + 16
180
+
181
+ def func_123(x):
182
+ return func_366(x) + 13
183
+
184
+ def func_124(x):
185
+ return func_706(x) - 2
186
+
187
+ def func_125(x):
188
+ return func_518(x) + 17
189
+
190
+ def func_126(x):
191
+ return func_1075(x) - 10
192
+
193
+ def func_130(x):
194
+ return func_960(x) - 12
195
+
196
+ def func_131(x):
197
+ return func_665(x) + 1
198
+
199
+ def func_132(x):
200
+ return func_650(x) + 13
201
+
202
+ def func_134(x):
203
+ return func_401(x) + 14
204
+
205
+ def func_137(x):
206
+ return func_979(x) - 6
207
+
208
+ def func_140(x):
209
+ return func_143(x) - 2
210
+
211
+ def func_141(x):
212
+ return func_599(x) - 11
213
+
214
+ def func_143(x):
215
+ return x + 3
216
+
217
+ def func_147(x):
218
+ return func_954(x) - 6
219
+
220
+ def func_153(x):
221
+ return func_371(x) + 3
222
+
223
+ def func_154(x):
224
+ return x + 3
225
+
226
+ def func_155(x):
227
+ return func_454(x)
228
+
229
+ def func_157(x):
230
+ return func_126(x) + 13
231
+
232
+ def func_158(x):
233
+ return func_319(x) + 10
234
+
235
+ def func_159(x):
236
+ return func_510(x) - 12
237
+
238
+ def func_162(x):
239
+ return func_707(x) + 8
240
+
241
+ def func_166(x):
242
+ return func_802(x) + 1
243
+
244
+ def func_167(x):
245
+ return func_1060(x) + 16
246
+
247
+ def func_169(x):
248
+ return func_741(x) - 11
249
+
250
+ def func_172(x):
251
+ return func_276(x) - 10
252
+
253
+ def func_176(x):
254
+ return func_23(x) + 1
255
+
256
+ def func_181(x):
257
+ return func_508(x) + 17
258
+
259
+ def func_185(x):
260
+ return func_1069(x) - 12
261
+
262
+ def func_188(x):
263
+ return func_1016(x) - 6
264
+
265
+ def func_194(x):
266
+ return func_661(x) - 1
267
+
268
+ def func_195(x):
269
+ return func_892(x) - 9
270
+
271
+ def func_198(x):
272
+ return x + 3
273
+
274
+ def func_199(x):
275
+ return func_716(x) + 3
276
+
277
+ def func_200(x):
278
+ return func_269(x) - 8
279
+
280
+ def func_201(x):
281
+ return func_943(x) + 14
282
+
283
+ def func_204(x):
284
+ return func_906(x) + 1
285
+
286
+ def func_205(x):
287
+ return func_1078(x) - 5
288
+
289
+ def func_207(x):
290
+ return func_167(x) - 4
291
+
292
+ def func_208(x):
293
+ return func_506(x) - 5
294
+
295
+ def func_209(x):
296
+ return func_1019(x)
297
+
298
+ def func_214(x):
299
+ return x + 9
300
+
301
+ def func_221(x):
302
+ return func_903(x) + 3
303
+
304
+ def func_224(x):
305
+ return x + 4
306
+
307
+ def func_225(x):
308
+ return func_480(x) + 6
309
+
310
+ def func_228(x):
311
+ return func_811(x) - 3
312
+
313
+ def func_231(x):
314
+ return func_490(x) + 16
315
+
316
+ def func_233(x):
317
+ return func_267(x) + 8
318
+
319
+ def func_234(x):
320
+ return func_541(x) + 8
321
+
322
+ def func_237(x):
323
+ return func_562(x)
324
+
325
+ def func_240(x):
326
+ return func_225(x) + 4
327
+
328
+ def func_242(x):
329
+ return func_432(x) + 8
330
+
331
+ def func_243(x):
332
+ return func_627(x) - 5
333
+
334
+ def func_244(x):
335
+ return func_23(x) + 9
336
+
337
+ def func_245(x):
338
+ return func_567(x) + 16
339
+
340
+ def func_248(x):
341
+ return func_115(x) + 5
342
+
343
+ def func_249(x):
344
+ return func_158(x) - 4
345
+
346
+ def func_251(x):
347
+ return x - 12
348
+
349
+ def func_253(x):
350
+ return func_403(x) - 12
351
+
352
+ def func_256(x):
353
+ return func_633(x) + 12
354
+
355
+ def func_262(x):
356
+ return func_917(x) - 12
357
+
358
+ def func_263(x):
359
+ return func_94(x) + 10
360
+
361
+ def func_265(x):
362
+ return func_1010(x) + 5
363
+
364
+ def func_267(x):
365
+ return func_681(x) + 11
366
+
367
+ def func_268(x):
368
+ return func_444(x) - 11
369
+
370
+ def func_269(x):
371
+ return func_717(x) + 13
372
+
373
+ def func_272(x):
374
+ return func_562(x) - 3
375
+
376
+ def func_274(x):
377
+ return func_820(x) + 15
378
+
379
+ def func_275(x):
380
+ return func_571(x) - 8
381
+
382
+ def func_276(x):
383
+ return x
384
+
385
+ def func_277(x):
386
+ return func_198(x) - 9
387
+
388
+ def func_278(x):
389
+ return func_1095(x) + 16
390
+
391
+ def func_279(x):
392
+ return func_525(x) + 3
393
+
394
+ def func_280(x):
395
+ return func_1029(x) - 12
396
+
397
+ def func_284(x):
398
+ return func_413(x) + 5
399
+
400
+ def func_286(x):
401
+ return x - 5
402
+
403
+ def func_287(x):
404
+ return func_101(x) - 7
405
+
406
+ def func_288(x):
407
+ return func_963(x) + 12
408
+
409
+ def func_289(x):
410
+ return func_16(x) + 15
411
+
412
+ def func_291(x):
413
+ return func_147(x) + 17
414
+
415
+ def func_292(x):
416
+ return func_405(x) + 12
417
+
418
+ def func_294(x):
419
+ return func_95(x)
420
+
421
+ def func_296(x):
422
+ return x + 17
423
+
424
+ def func_297(x):
425
+ return func_140(x) + 11
426
+
427
+ def func_299(x):
428
+ return func_274(x) + 10
429
+
430
+ def func_301(x):
431
+ return func_113(x) + 9
432
+
433
+ def func_302(x):
434
+ return func_1086(x) - 9
435
+
436
+ def func_303(x):
437
+ return func_521(x) + 17
438
+
439
+ def func_308(x):
440
+ return func_727(x) - 11
441
+
442
+ def func_309(x):
443
+ return func_302(x) + 5
444
+
445
+ def func_310(x):
446
+ return func_48(x) - 12
447
+
448
+ def func_313(x):
449
+ return x + 6
450
+
451
+ def func_315(x):
452
+ return x - 5
453
+
454
+ def func_316(x):
455
+ return func_670(x) + 12
456
+
457
+ def func_317(x):
458
+ return func_1005(x) + 15
459
+
460
+ def func_319(x):
461
+ return func_98(x) - 4
462
+
463
+ def func_320(x):
464
+ return x + 5
465
+
466
+ def func_323(x):
467
+ return func_657(x) - 4
468
+
469
+ def func_324(x):
470
+ return func_877(x) - 9
471
+
472
+ def func_325(x):
473
+ return func_320(x) - 5
474
+
475
+ def func_327(x):
476
+ return func_757(x) - 9
477
+
478
+ def func_330(x):
479
+ return func_825(x) - 4
480
+
481
+ def func_334(x):
482
+ return func_122(x)
483
+
484
+ def func_335(x):
485
+ return func_445(x) - 7
486
+
487
+ def func_336(x):
488
+ return func_153(x) + 16
489
+
490
+ def func_340(x):
491
+ return func_758(x) - 10
492
+
493
+ def func_346(x):
494
+ return func_85(x) + 1
495
+
496
+ def func_348(x):
497
+ return func_567(x) + 8
498
+
499
+ def func_351(x):
500
+ return func_22(x) + 5
501
+
502
+ def func_352(x):
503
+ return func_527(x) + 16
504
+
505
+ def func_353(x):
506
+ return func_860(x) - 7
507
+
508
+ def func_357(x):
509
+ return func_878(x) + 1
510
+
511
+ def func_358(x):
512
+ return func_960(x) - 11
513
+
514
+ def func_361(x):
515
+ return func_48(x) + 5
516
+
517
+ def func_362(x):
518
+ return func_134(x) - 2
519
+
520
+ def func_363(x):
521
+ return func_1095(x) - 5
522
+
523
+ def func_364(x):
524
+ return func_346(x) - 7
525
+
526
+ def func_365(x):
527
+ return func_527(x) - 7
528
+
529
+ def func_366(x):
530
+ return func_361(x) - 1
531
+
532
+ def func_367(x):
533
+ return func_375(x) + 17
534
+
535
+ def func_369(x):
536
+ return x - 5
537
+
538
+ def func_370(x):
539
+ return func_556(x) + 1
540
+
541
+ def func_371(x):
542
+ return func_141(x) - 10
543
+
544
+ def func_374(x):
545
+ return x - 2
546
+
547
+ def func_375(x):
548
+ return func_828(x) - 6
549
+
550
+ def func_376(x):
551
+ return func_251(x) - 5
552
+
553
+ def func_378(x):
554
+ return func_231(x) - 8
555
+
556
+ def func_382(x):
557
+ return func_1080(x) - 8
558
+
559
+ def func_385(x):
560
+ return func_1067(x) + 11
561
+
562
+ def func_386(x):
563
+ return func_1003(x) + 14
564
+
565
+ def func_387(x):
566
+ return func_98(x) - 9
567
+
568
+ def func_390(x):
569
+ return func_1062(x) + 15
570
+
571
+ def func_391(x):
572
+ return func_486(x) + 5
573
+
574
+ def func_392(x):
575
+ return func_88(x) - 1
576
+
577
+ def func_393(x):
578
+ return func_3(x) + 3
579
+
580
+ def func_395(x):
581
+ return func_741(x)
582
+
583
+ def func_397(x):
584
+ return func_730(x) + 17
585
+
586
+ def func_400(x):
587
+ return func_253(x) + 1
588
+
589
+ def func_401(x):
590
+ return func_376(x) + 10
591
+
592
+ def func_402(x):
593
+ return func_556(x) + 9
594
+
595
+ def func_403(x):
596
+ return func_506(x) + 13
597
+
598
+ def func_405(x):
599
+ return func_572(x) + 13
600
+
601
+ def func_406(x):
602
+ return x + 3
603
+
604
+ def func_413(x):
605
+ return func_90(x) - 9
606
+
607
+ def func_422(x):
608
+ return func_770(x) + 17
609
+
610
+ def func_423(x):
611
+ return func_1049(x) - 10
612
+
613
+ def func_428(x):
614
+ return func_278(x) + 12
615
+
616
+ def func_429(x):
617
+ return func_931(x) - 8
618
+
619
+ def func_432(x):
620
+ return func_292(x) - 8
621
+
622
+ def func_434(x):
623
+ return x + 2
624
+
625
+ def func_441(x):
626
+ return func_297(x) + 11
627
+
628
+ def func_443(x):
629
+ return func_696(x) + 12
630
+
631
+ def func_444(x):
632
+ return func_124(x) + 16
633
+
634
+ def func_445(x):
635
+ return func_618(x) - 5
636
+
637
+ def func_454(x):
638
+ return func_113(x) - 4
639
+
640
+ def func_455(x):
641
+ return func_325(x) - 2
642
+
643
+ def func_456(x):
644
+ return func_1007(x) + 7
645
+
646
+ def func_457(x):
647
+ return func_284(x) - 11
648
+
649
+ def func_458(x):
650
+ return func_789(x) + 1
651
+
652
+ def func_461(x):
653
+ return func_859(x) + 16
654
+
655
+ def func_462(x):
656
+ return func_1083(x) - 6
657
+
658
+ def func_463(x):
659
+ return func_456(x) + 11
660
+
661
+ def func_466(x):
662
+ return func_403(x) - 1
663
+
664
+ def func_469(x):
665
+ return func_698(x) + 13
666
+
667
+ def func_470(x):
668
+ return func_251(x) + 7
669
+
670
+ def func_473(x):
671
+ return func_910(x) + 5
672
+
673
+ def func_474(x):
674
+ return func_688(x) + 10
675
+
676
+ def func_480(x):
677
+ return func_1007(x) - 7
678
+
679
+ def func_484(x):
680
+ return func_673(x) + 3
681
+
682
+ def func_486(x):
683
+ return func_12(x) + 2
684
+
685
+ def func_487(x):
686
+ return func_70(x) - 11
687
+
688
+ def func_490(x):
689
+ return func_455(x) - 2
690
+
691
+ def func_492(x):
692
+ return func_53(x) + 7
693
+
694
+ def func_493(x):
695
+ return func_288(x) - 8
696
+
697
+ def func_494(x):
698
+ return func_757(x) + 4
699
+
700
+ def func_495(x):
701
+ return x - 11
702
+
703
+ def func_503(x):
704
+ return func_801(x) + 4
705
+
706
+ def func_504(x):
707
+ return func_1005(x) - 5
708
+
709
+ def func_505(x):
710
+ return func_102(x) - 11
711
+
712
+ def func_506(x):
713
+ return func_865(x) + 16
714
+
715
+ def func_508(x):
716
+ return func_863(x) + 13
717
+
718
+ def func_510(x):
719
+ return func_348(x) - 3
720
+
721
+ def func_514(x):
722
+ return func_302(x) - 4
723
+
724
+ def func_516(x):
725
+ return func_558(x) + 9
726
+
727
+ def func_518(x):
728
+ return func_36(x) + 11
729
+
730
+ def func_521(x):
731
+ return func_658(x) + 1
732
+
733
+ def func_523(x):
734
+ return func_960(x) - 8
735
+
736
+ def func_525(x):
737
+ return func_95(x) + 14
738
+
739
+ def func_526(x):
740
+ return func_249(x) - 4
741
+
742
+ def func_527(x):
743
+ return x + 8
744
+
745
+ def func_529(x):
746
+ return func_627(x) + 17
747
+
748
+ def func_531(x):
749
+ return func_323(x) + 14
750
+
751
+ def func_532(x):
752
+ return func_1010(x) + 6
753
+
754
+ def func_533(x):
755
+ return func_158(x) - 8
756
+
757
+ def func_538(x):
758
+ return func_864(x) + 10
759
+
760
+ def func_540(x):
761
+ return func_121(x) - 12
762
+
763
+ def func_541(x):
764
+ return func_131(x) - 10
765
+
766
+ def func_542(x):
767
+ return func_1077(x) + 12
768
+
769
+ def func_543(x):
770
+ return func_233(x) + 8
771
+
772
+ def func_545(x):
773
+ return func_240(x) + 5
774
+
775
+ def func_547(x):
776
+ return func_126(x) + 9
777
+
778
+ def func_548(x):
779
+ return x + 6
780
+
781
+ def func_549(x):
782
+ return func_395(x) - 8
783
+
784
+ def func_550(x):
785
+ return func_650(x) - 5
786
+
787
+ def func_552(x):
788
+ return func_324(x) - 5
789
+
790
+ def func_555(x):
791
+ return x - 10
792
+
793
+ def func_556(x):
794
+ return func_1089(x)
795
+
796
+ def func_557(x):
797
+ return func_32(x) + 17
798
+
799
+ def func_558(x):
800
+ return func_952(x) - 9
801
+
802
+ def func_559(x):
803
+ return func_397(x) + 15
804
+
805
+ def func_561(x):
806
+ return func_1031(x) + 17
807
+
808
+ def func_562(x):
809
+ return func_71(x) - 4
810
+
811
+ def func_564(x):
812
+ return func_1095(x) + 4
813
+
814
+ def func_565(x):
815
+ return func_432(x) - 7
816
+
817
+ def func_567(x):
818
+ return func_778(x) - 5
819
+
820
+ def func_571(x):
821
+ return func_552(x) + 2
822
+
823
+ def func_572(x):
824
+ return func_251(x) - 8
825
+
826
+ def func_576(x):
827
+ return func_251(x) - 1
828
+
829
+ def func_578(x):
830
+ return func_860(x) - 12
831
+
832
+ def func_579(x):
833
+ return func_141(x) + 16
834
+
835
+ def func_584(x):
836
+ return func_249(x) + 16
837
+
838
+ def func_588(x):
839
+ return func_1020(x) + 13
840
+
841
+ def func_590(x):
842
+ return func_382(x) - 9
843
+
844
+ def func_594(x):
845
+ return func_262(x) - 10
846
+
847
+ def func_595(x):
848
+ return func_662(x) + 5
849
+
850
+ def func_596(x):
851
+ return func_275(x) + 9
852
+
853
+ def func_599(x):
854
+ return x + 6
855
+
856
+ def func_600(x):
857
+ return func_699(x) + 7
858
+
859
+ def func_610(x):
860
+ return x - 1
861
+
862
+ def func_611(x):
863
+ return func_169(x) + 3
864
+
865
+ def func_612(x):
866
+ return func_979(x) + 6
867
+
868
+ def func_617(x):
869
+ return func_875(x) + 7
870
+
871
+ def func_618(x):
872
+ return func_313(x) - 2
873
+
874
+ def func_620(x):
875
+ return func_796(x) + 9
876
+
877
+ def func_622(x):
878
+ return func_1089(x) - 7
879
+
880
+ def func_625(x):
881
+ return func_101(x) - 12
882
+
883
+ def func_626(x):
884
+ return func_474(x) - 10
885
+
886
+ def func_627(x):
887
+ return func_1060(x) - 5
888
+
889
+ def func_633(x):
890
+ return func_879(x) - 8
891
+
892
+ def func_634(x):
893
+ return func_292(x) + 2
894
+
895
+ def func_637(x):
896
+ return func_25(x) + 7
897
+
898
+ def func_638(x):
899
+ return func_36(x) - 3
900
+
901
+ def func_639(x):
902
+ return func_316(x) + 12
903
+
904
+ def func_641(x):
905
+ return func_829(x) - 9
906
+
907
+ def func_644(x):
908
+ return func_662(x) - 11
909
+
910
+ def func_645(x):
911
+ return func_965(x) + 9
912
+
913
+ def func_647(x):
914
+ return func_1007(x) - 10
915
+
916
+ def func_648(x):
917
+ return func_548(x) + 1
918
+
919
+ def func_649(x):
920
+ return func_692(x) + 13
921
+
922
+ def func_650(x):
923
+ return func_1010(x)
924
+
925
+ def func_653(x):
926
+ return func_1086(x) - 12
927
+
928
+ def func_657(x):
929
+ return func_90(x) + 4
930
+
931
+ def func_658(x):
932
+ return func_761(x) - 5
933
+
934
+ def func_659(x):
935
+ return func_14(x) - 2
936
+
937
+ def func_661(x):
938
+ return func_853(x) - 12
939
+
940
+ def func_662(x):
941
+ return func_872(x) + 16
942
+
943
+ def func_664(x):
944
+ return func_245(x) + 7
945
+
946
+ def func_665(x):
947
+ return func_251(x) + 5
948
+
949
+ def func_669(x):
950
+ return func_657(x) + 2
951
+
952
+ def func_670(x):
953
+ return x + 11
954
+
955
+ def func_672(x):
956
+ return x - 4
957
+
958
+ def func_673(x):
959
+ return func_869(x) - 4
960
+
961
+ def func_674(x):
962
+ return x - 8
963
+
964
+ def func_675(x):
965
+ return func_291(x) - 12
966
+
967
+ def func_676(x):
968
+ return func_599(x) + 10
969
+
970
+ def func_677(x):
971
+ return func_423(x) + 17
972
+
973
+ def func_678(x):
974
+ return func_758(x) + 7
975
+
976
+ def func_679(x):
977
+ return func_119(x) + 17
978
+
979
+ def func_681(x):
980
+ return x - 7
981
+
982
+ def func_683(x):
983
+ return func_1029(x) + 3
984
+
985
+ def func_685(x):
986
+ return func_248(x) + 11
987
+
988
+ def func_686(x):
989
+ return func_1099(x) + 7
990
+
991
+ def func_688(x):
992
+ return func_910(x) + 3
993
+
994
+ def func_692(x):
995
+ return func_997(x) + 7
996
+
997
+ def func_693(x):
998
+ return func_391(x) - 11
999
+
1000
+ def func_694(x):
1001
+ return x + 5
1002
+
1003
+ def func_695(x):
1004
+ return func_262(x) + 6
1005
+
1006
+ def func_696(x):
1007
+ return func_1013(x) - 5
1008
+
1009
+ def func_698(x):
1010
+ return func_890(x) + 5
1011
+
1012
+ def func_699(x):
1013
+ return func_965(x)
1014
+
1015
+ def func_701(x):
1016
+ return func_386(x) + 15
1017
+
1018
+ def func_702(x):
1019
+ return func_30(x) + 16
1020
+
1021
+ def func_703(x):
1022
+ return func_1007(x) - 6
1023
+
1024
+ def func_705(x):
1025
+ return func_964(x) - 1
1026
+
1027
+ def func_706(x):
1028
+ return func_308(x) + 14
1029
+
1030
+ def func_707(x):
1031
+ return x - 8
1032
+
1033
+ def func_715(x):
1034
+ return func_826(x) - 6
1035
+
1036
+ def func_716(x):
1037
+ return func_741(x) - 6
1038
+
1039
+ def func_717(x):
1040
+ return func_454(x) - 5
1041
+
1042
+ def func_718(x):
1043
+ return func_242(x)
1044
+
1045
+ def func_721(x):
1046
+ return x + 9
1047
+
1048
+ def func_722(x):
1049
+ return func_14(x) - 11
1050
+
1051
+ def func_723(x):
1052
+ return func_693(x) - 4
1053
+
1054
+ def func_727(x):
1055
+ return func_647(x) + 13
1056
+
1057
+ def func_729(x):
1058
+ return func_989(x) - 9
1059
+
1060
+ def func_730(x):
1061
+ return func_617(x) + 1
1062
+
1063
+ def func_731(x):
1064
+ return func_124(x) + 17
1065
+
1066
+ def func_732(x):
1067
+ return func_443(x) + 12
1068
+
1069
+ def func_735(x):
1070
+ return func_253(x) + 6
1071
+
1072
+ def func_739(x):
1073
+ return func_829(x)
1074
+
1075
+ def func_740(x):
1076
+ return func_369(x) + 12
1077
+
1078
+ def func_741(x):
1079
+ return x + 1
1080
+
1081
+ def func_746(x):
1082
+ return func_267(x) + 6
1083
+
1084
+ def func_747(x):
1085
+ return func_699(x) + 4
1086
+
1087
+ def func_750(x):
1088
+ return func_527(x) + 7
1089
+
1090
+ def func_751(x):
1091
+ return func_1067(x) + 8
1092
+
1093
+ def func_754(x):
1094
+ return func_960(x) + 17
1095
+
1096
+ def func_756(x):
1097
+ return x + 14
1098
+
1099
+ def func_757(x):
1100
+ return func_58(x) - 5
1101
+
1102
+ def func_758(x):
1103
+ return func_1078(x) + 13
1104
+
1105
+ def func_759(x):
1106
+ return func_70(x) + 9
1107
+
1108
+ def func_760(x):
1109
+ return func_943(x) - 4
1110
+
1111
+ def func_761(x):
1112
+ return func_325(x) + 4
1113
+
1114
+ def func_768(x):
1115
+ return func_637(x)
1116
+
1117
+ def func_769(x):
1118
+ return func_692(x) - 9
1119
+
1120
+ def func_770(x):
1121
+ return func_679(x) - 12
1122
+
1123
+ def func_772(x):
1124
+ return func_1016(x)
1125
+
1126
+ def func_778(x):
1127
+ return func_224(x) - 11
1128
+
1129
+ def func_782(x):
1130
+ return func_118(x) + 4
1131
+
1132
+ def func_787(x):
1133
+ return x - 9
1134
+
1135
+ def func_789(x):
1136
+ return x + 10
1137
+
1138
+ def func_792(x):
1139
+ return x + 4
1140
+
1141
+ def func_796(x):
1142
+ return func_770(x) - 7
1143
+
1144
+ def func_798(x):
1145
+ return func_1044(x) + 14
1146
+
1147
+ def func_800(x):
1148
+ return func_527(x) + 14
1149
+
1150
+ def func_801(x):
1151
+ return func_971(x) - 7
1152
+
1153
+ def func_802(x):
1154
+ return func_92(x) - 9
1155
+
1156
+ def func_804(x):
1157
+ return func_70(x) + 2
1158
+
1159
+ def func_805(x):
1160
+ return func_676(x) - 2
1161
+
1162
+ def func_811(x):
1163
+ return func_741(x) + 9
1164
+
1165
+ def func_812(x):
1166
+ return func_176(x) + 17
1167
+
1168
+ def func_813(x):
1169
+ return func_114(x) - 3
1170
+
1171
+ def func_814(x):
1172
+ return func_851(x) + 10
1173
+
1174
+ def func_815(x):
1175
+ return func_361(x) + 13
1176
+
1177
+ def func_819(x):
1178
+ return func_730(x) + 9
1179
+
1180
+ def func_820(x):
1181
+ return func_248(x) - 11
1182
+
1183
+ def func_821(x):
1184
+ return func_233(x) - 10
1185
+
1186
+ def func_823(x):
1187
+ return func_819(x) - 3
1188
+
1189
+ def func_824(x):
1190
+ return func_622(x) + 5
1191
+
1192
+ def func_825(x):
1193
+ return func_176(x) + 15
1194
+
1195
+ def func_826(x):
1196
+ return func_1047(x) - 5
1197
+
1198
+ def func_827(x):
1199
+ return func_625(x) + 3
1200
+
1201
+ def func_828(x):
1202
+ return func_126(x) - 10
1203
+
1204
+ def func_829(x):
1205
+ return func_815(x) + 12
1206
+
1207
+ def func_830(x):
1208
+ return func_863(x) + 3
1209
+
1210
+ def func_832(x):
1211
+ return func_401(x) - 11
1212
+
1213
+ def func_836(x):
1214
+ return func_492(x) + 12
1215
+
1216
+ def func_838(x):
1217
+ return func_153(x) + 14
1218
+
1219
+ def func_840(x):
1220
+ return x - 3
1221
+
1222
+ def func_842(x):
1223
+ return func_253(x) - 3
1224
+
1225
+ def func_843(x):
1226
+ return func_987(x) + 1
1227
+
1228
+ def func_845(x):
1229
+ return func_463(x) - 7
1230
+
1231
+ def func_846(x):
1232
+ return func_678(x) + 3
1233
+
1234
+ def func_847(x):
1235
+ return func_199(x) - 6
1236
+
1237
+ def func_851(x):
1238
+ return func_505(x) - 4
1239
+
1240
+ def func_853(x):
1241
+ return func_990(x) + 8
1242
+
1243
+ def func_856(x):
1244
+ return func_397(x) + 16
1245
+
1246
+ def func_857(x):
1247
+ return func_579(x) - 3
1248
+
1249
+ def func_859(x):
1250
+ return func_406(x) + 1
1251
+
1252
+ def func_860(x):
1253
+ return func_378(x) + 14
1254
+
1255
+ def func_861(x):
1256
+ return func_958(x)
1257
+
1258
+ def func_863(x):
1259
+ return func_361(x) - 4
1260
+
1261
+ def func_864(x):
1262
+ return func_730(x) + 2
1263
+
1264
+ def func_865(x):
1265
+ return x - 6
1266
+
1267
+ def func_866(x):
1268
+ return x + 4
1269
+
1270
+ def func_869(x):
1271
+ return func_369(x) + 1
1272
+
1273
+ def func_871(x):
1274
+ return func_265(x) + 3
1275
+
1276
+ def func_872(x):
1277
+ return func_902(x) + 17
1278
+
1279
+ def func_873(x):
1280
+ return func_1076(x) + 14
1281
+
1282
+ def func_875(x):
1283
+ return func_309(x) + 1
1284
+
1285
+ def func_877(x):
1286
+ return func_750(x) + 9
1287
+
1288
+ def func_878(x):
1289
+ return func_1021(x) - 11
1290
+
1291
+ def func_879(x):
1292
+ return func_423(x) + 16
1293
+
1294
+ def func_880(x):
1295
+ return func_1042(x) + 7
1296
+
1297
+ def func_882(x):
1298
+ return func_527(x) - 1
1299
+
1300
+ def func_886(x):
1301
+ return func_1091(x)
1302
+
1303
+ def func_887(x):
1304
+ return func_208(x) + 12
1305
+
1306
+ def func_889(x):
1307
+ return func_36(x) - 11
1308
+
1309
+ def func_890(x):
1310
+ return func_1091(x) - 8
1311
+
1312
+ def func_891(x):
1313
+ return func_492(x) + 14
1314
+
1315
+ def func_892(x):
1316
+ return func_233(x) + 16
1317
+
1318
+ def func_896(x):
1319
+ return func_827(x) + 7
1320
+
1321
+ def func_901(x):
1322
+ return func_284(x) + 11
1323
+
1324
+ def func_902(x):
1325
+ return func_406(x) + 5
1326
+
1327
+ def func_903(x):
1328
+ return func_23(x) + 2
1329
+
1330
+ def func_906(x):
1331
+ return func_301(x) - 1
1332
+
1333
+ def func_907(x):
1334
+ return func_578(x) + 2
1335
+
1336
+ def func_910(x):
1337
+ return func_195(x) - 9
1338
+
1339
+ def func_911(x):
1340
+ return func_983(x) + 7
1341
+
1342
+ def func_912(x):
1343
+ return x + 15
1344
+
1345
+ def func_913(x):
1346
+ return x - 6
1347
+
1348
+ def func_915(x):
1349
+ return func_1080(x) - 2
1350
+
1351
+ def func_917(x):
1352
+ return func_693(x) - 7
1353
+
1354
+ def func_920(x):
1355
+ return func_516(x) + 16
1356
+
1357
+ def func_923(x):
1358
+ return func_336(x) - 1
1359
+
1360
+ def func_924(x):
1361
+ return func_443(x) - 12
1362
+
1363
+ def func_927(x):
1364
+ return func_7(x) + 15
1365
+
1366
+ def func_928(x):
1367
+ return func_335(x) + 2
1368
+
1369
+ def func_931(x):
1370
+ return func_245(x)
1371
+
1372
+ def func_934(x):
1373
+ return func_1042(x) - 1
1374
+
1375
+ def func_936(x):
1376
+ return func_137(x) + 6
1377
+
1378
+ def func_937(x):
1379
+ return func_915(x) + 4
1380
+
1381
+ def func_939(x):
1382
+ return func_353(x) + 14
1383
+
1384
+ def func_940(x):
1385
+ return func_757(x) - 7
1386
+
1387
+ def func_943(x):
1388
+ return func_208(x) + 14
1389
+
1390
+ def func_945(x):
1391
+ return func_330(x) + 5
1392
+
1393
+ def func_948(x):
1394
+ return func_686(x) - 11
1395
+
1396
+ def func_949(x):
1397
+ return func_757(x) + 13
1398
+
1399
+ def func_950(x):
1400
+ return x + 5
1401
+
1402
+ def func_952(x):
1403
+ return func_493(x) + 13
1404
+
1405
+ def func_953(x):
1406
+ return x + 17
1407
+
1408
+ def func_954(x):
1409
+ return x - 7
1410
+
1411
+ def func_955(x):
1412
+ return func_772(x) + 2
1413
+
1414
+ def func_957(x):
1415
+ return func_948(x)
1416
+
1417
+ def func_958(x):
1418
+ return func_578(x) - 10
1419
+
1420
+ def func_960(x):
1421
+ return func_677(x) - 6
1422
+
1423
+ def func_962(x):
1424
+ return func_564(x) + 11
1425
+
1426
+ def func_963(x):
1427
+ return func_1007(x) - 5
1428
+
1429
+ def func_964(x):
1430
+ return func_286(x) + 9
1431
+
1432
+ def func_965(x):
1433
+ return func_375(x) + 7
1434
+
1435
+ def func_971(x):
1436
+ return func_953(x) - 10
1437
+
1438
+ def func_972(x):
1439
+ return func_564(x) - 12
1440
+
1441
+ def func_973(x):
1442
+ return x + 11
1443
+
1444
+ def func_974(x):
1445
+ return func_637(x) + 3
1446
+
1447
+ def func_976(x):
1448
+ return func_696(x) - 6
1449
+
1450
+ def func_978(x):
1451
+ return func_461(x) - 4
1452
+
1453
+ def func_979(x):
1454
+ return func_672(x) - 9
1455
+
1456
+ def func_983(x):
1457
+ return func_648(x) + 4
1458
+
1459
+ def func_985(x):
1460
+ return func_564(x) - 10
1461
+
1462
+ def func_986(x):
1463
+ return func_936(x) - 5
1464
+
1465
+ def func_987(x):
1466
+ return func_873(x) + 3
1467
+
1468
+ def func_988(x):
1469
+ return x + 7
1470
+
1471
+ def func_989(x):
1472
+ return func_335(x) + 8
1473
+
1474
+ def func_990(x):
1475
+ return func_674(x) - 9
1476
+
1477
+ def func_991(x):
1478
+ return func_1067(x) + 1
1479
+
1480
+ def func_992(x):
1481
+ return func_351(x)
1482
+
1483
+ def func_993(x):
1484
+ return func_1043(x) + 7
1485
+
1486
+ def func_996(x):
1487
+ return func_896(x) + 13
1488
+
1489
+ def func_997(x):
1490
+ return func_688(x) - 6
1491
+
1492
+ def func_1000(x):
1493
+ return func_986(x) + 5
1494
+
1495
+ def func_1003(x):
1496
+ return func_296(x) - 6
1497
+
1498
+ def func_1004(x):
1499
+ return func_463(x) - 1
1500
+
1501
+ def func_1005(x):
1502
+ return func_92(x) + 1
1503
+
1504
+ def func_1007(x):
1505
+ return func_572(x) - 1
1506
+
1507
+ def func_1008(x):
1508
+ return func_367(x) + 17
1509
+
1510
+ def func_1010(x):
1511
+ return func_224(x) - 12
1512
+
1513
+ def func_1013(x):
1514
+ return func_262(x) + 15
1515
+
1516
+ def func_1016(x):
1517
+ return func_276(x) + 1
1518
+
1519
+ def func_1019(x):
1520
+ return x - 10
1521
+
1522
+ def func_1020(x):
1523
+ return func_782(x) + 8
1524
+
1525
+ def func_1021(x):
1526
+ return x + 12
1527
+
1528
+ def func_1027(x):
1529
+ return func_405(x) + 2
1530
+
1531
+ def func_1029(x):
1532
+ return func_221(x) + 3
1533
+
1534
+ def func_1030(x):
1535
+ return func_237(x) - 8
1536
+
1537
+ def func_1031(x):
1538
+ return func_12(x) - 2
1539
+
1540
+ def func_1032(x):
1541
+ return func_813(x) + 16
1542
+
1543
+ def func_1035(x):
1544
+ return func_294(x) + 5
1545
+
1546
+ def func_1037(x):
1547
+ return func_954(x) + 17
1548
+
1549
+ def func_1042(x):
1550
+ return func_23(x) + 11
1551
+
1552
+ def func_1043(x):
1553
+ return func_845(x) + 6
1554
+
1555
+ def func_1044(x):
1556
+ return x - 7
1557
+
1558
+ def func_1045(x):
1559
+ return x + 11
1560
+
1561
+ def func_1047(x):
1562
+ return func_288(x) + 1
1563
+
1564
+ def func_1049(x):
1565
+ return func_88(x) - 6
1566
+
1567
+ def func_1051(x):
1568
+ return func_63(x) - 4
1569
+
1570
+ def func_1053(x):
1571
+ return func_832(x) - 5
1572
+
1573
+ def func_1054(x):
1574
+ return func_761(x) - 3
1575
+
1576
+ def func_1059(x):
1577
+ return func_397(x) + 12
1578
+
1579
+ def func_1060(x):
1580
+ return func_600(x) + 17
1581
+
1582
+ def func_1061(x):
1583
+ return func_826(x) + 6
1584
+
1585
+ def func_1062(x):
1586
+ return func_549(x) + 4
1587
+
1588
+ def func_1067(x):
1589
+ return func_963(x) + 2
1590
+
1591
+ def func_1069(x):
1592
+ return func_541(x) + 7
1593
+
1594
+ def func_1075(x):
1595
+ return x + 7
1596
+
1597
+ def func_1076(x):
1598
+ return func_845(x) + 11
1599
+
1600
+ def func_1077(x):
1601
+ return func_661(x) - 10
1602
+
1603
+ def func_1078(x):
1604
+ return func_634(x) - 7
1605
+
1606
+ def func_1079(x):
1607
+ return func_928(x) - 11
1608
+
1609
+ def func_1080(x):
1610
+ return func_658(x) + 6
1611
+
1612
+ def func_1082(x):
1613
+ return x + 6
1614
+
1615
+ def func_1083(x):
1616
+ return func_237(x) + 4
1617
+
1618
+ def func_1086(x):
1619
+ return func_1082(x) - 3
1620
+
1621
+ def func_1089(x):
1622
+ return func_625(x) + 14
1623
+
1624
+ def func_1090(x):
1625
+ return func_760(x) - 10
1626
+
1627
+ def func_1091(x):
1628
+ return func_393(x) + 13
1629
+
1630
+ def func_1093(x):
1631
+ return func_244(x) - 5
1632
+
1633
+ def func_1094(x):
1634
+ return func_813(x) - 9
1635
+
1636
+ def func_1095(x):
1637
+ return func_387(x) - 8
1638
+
1639
+ def func_1096(x):
1640
+ return func_185(x) - 8
1641
+
1642
+ def func_1098(x):
1643
+ return func_873(x) + 1
1644
+
1645
+ def func_1099(x):
1646
+ return func_456(x) - 8
1647
+
1648
+ def func_1100(x):
1649
+ return func_692(x)
1650
+
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ openai
2
+ tiktoken
3
+ rouge
4
+ torch
5
+ transformers
6
+ accelerate
7
+ evaluate
8
+ xopen
9
+ python-dotenv
App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/scripts/download_dataset.sh ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ save_dir=data
3
+ mkdir ${save_dir}
4
+ for file in code_debug code_run kv_retrieval longbook_choice_eng longbook_qa_chn longbook_qa_eng longbook_sum_eng longdialogue_qa_eng math_calc math_find number_string passkey; do
5
+ wget -c https://huggingface.co/datasets/xinrongzhang2022/InfiniteBench/resolve/main/${file}.jsonl?download=true -O ./${save_dir}/${file}.jsonl
6
+ done
App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/MMLU_Pro_rewritten.py ADDED
@@ -0,0 +1,341 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # MMLU_Pro_rewritten.py
2
+ # Description: Script to perform MMLU-Pro benchmarking
3
+ #
4
+ ####################################################################################################################
5
+ # Imports
6
+ import os
7
+ import threading
8
+ import time
9
+ import toml
10
+ from tqdm import tqdm
11
+ from concurrent.futures import ThreadPoolExecutor
12
+ import logging
13
+ from openai import OpenAI
14
+ from datasets import load_dataset
15
+ import json
16
+ import re
17
+ #
18
+ ##################################################################################################################
19
+ #
20
+ # Functions:
21
+
22
+
23
+ # Set up logging
24
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
25
+ logger = logging.getLogger(__name__)
26
+
27
+
28
+ def load_mmlu_pro_config(**kwargs):
29
+ # Get the directory of the current script
30
+ script_dir = os.path.dirname(os.path.abspath(__file__))
31
+
32
+ # Construct the full path to config.toml
33
+ config_path = os.path.join(script_dir, 'config.toml')
34
+
35
+ # Load the config
36
+ config = toml.load(config_path)
37
+
38
+ # Update config with provided kwargs
39
+ for key, value in kwargs.items():
40
+ if key in config["server"]:
41
+ config["server"][key] = value
42
+ elif key in config["test"]:
43
+ config["test"][key] = value
44
+ elif key in config["log"]:
45
+ config["log"][key] = value
46
+
47
+ return config
48
+
49
+ # client_initializer.py
50
+ def initialize_client(config):
51
+ try:
52
+ return OpenAI(
53
+ base_url=config["server"]["url"],
54
+ api_key=config["server"]["api_key"],
55
+ timeout=config["server"]["timeout"]
56
+ )
57
+ except Exception as e:
58
+ logger.error(f"Failed to initialize OpenAI client: {e}")
59
+ raise
60
+
61
+ # dataset_loader.py
62
+ def load_mmlu_pro():
63
+ try:
64
+ dataset = load_dataset("TIGER-Lab/MMLU-Pro")
65
+ test_df, val_df = dataset["test"], dataset["validation"]
66
+ return preprocess(test_df), preprocess(val_df)
67
+ except Exception as e:
68
+ logger.error(f"Error loading MMLU-Pro dataset: {e}")
69
+ raise
70
+
71
+ def preprocess(data):
72
+ res = {}
73
+ for item in data:
74
+ options = [opt for opt in item["options"] if opt != "N/A"]
75
+ item["options"] = options
76
+ category = item["category"]
77
+ if category not in res:
78
+ res[category] = []
79
+ res[category].append(item)
80
+ return res
81
+
82
+ # prompt_creator.py
83
+ def create_prompt(cot_examples, question, options, config):
84
+ style = config["inference"]["style"]
85
+ system_prompt = config["inference"]["system_prompt"]
86
+
87
+ def format_example(q, opts, cot=""):
88
+ if not cot:
89
+ cot = "Let's think step by step."
90
+ cot = cot[3:] if cot.startswith("A: ") else cot
91
+ example = f"Question: {q}\nOptions: "
92
+ example += "\n".join(f"{chr(65 + i)}. {opt}" for i, opt in enumerate(opts))
93
+ return example.strip(), cot.strip()
94
+
95
+ if style == "multi_chat":
96
+ messages = [{"role": "system", "content": system_prompt}]
97
+ for ex in cot_examples:
98
+ ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"])
99
+ messages.extend([
100
+ {"role": "user", "content": ex_text},
101
+ {"role": "assistant", "content": f"Answer: {cot}"}
102
+ ])
103
+ q_text, _ = format_example(question, options)
104
+ messages.append({"role": "user", "content": q_text})
105
+ return messages
106
+ elif style == "single_chat":
107
+ prompt = f"{system_prompt}\n\n"
108
+ for ex in cot_examples:
109
+ ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"])
110
+ prompt += f"{ex_text}\nAnswer: {cot}\n\n"
111
+ q_text, _ = format_example(question, options)
112
+ prompt += f"{q_text}\nAnswer: Let's think step by step."
113
+ return [{"role": "user", "content": prompt}]
114
+ else: # no_chat
115
+ prompt = f"{system_prompt}\n\n"
116
+ for ex in cot_examples:
117
+ ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"])
118
+ prompt += f"{ex_text}\nAnswer: {cot}\n\n"
119
+ q_text, _ = format_example(question, options)
120
+ prompt += f"{q_text}\nAnswer: Let's think step by step."
121
+ return prompt
122
+
123
+ # answer_extractor.py
124
+ def extract_answer(text):
125
+ patterns = [
126
+ r"answer is \(?([A-J])\)?",
127
+ r".*[aA]nswer:\s*\(?([A-J])\)?",
128
+ r"\b([A-J])\b(?!.*\b[A-J]\b)"
129
+ ]
130
+
131
+ for pattern in patterns:
132
+ match = re.search(pattern, text, re.IGNORECASE | re.DOTALL)
133
+ if match:
134
+ return match.group(1).upper()
135
+
136
+ logger.warning(f"Failed to extract answer from: {text}")
137
+ return None
138
+
139
+ # question_evaluator.py
140
+ def run_single_question(question, cot_examples, client, config):
141
+ max_retries = 3
142
+ for attempt in range(max_retries):
143
+ try:
144
+ prompt = create_prompt(cot_examples, question['question'], question['options'], config)
145
+
146
+ if config["inference"]["style"] == "no_chat":
147
+ response = client.completions.create(
148
+ model=config["server"]["model"],
149
+ prompt=prompt,
150
+ temperature=config["inference"]["temperature"],
151
+ max_tokens=config["inference"]["max_tokens"],
152
+ top_p=config["inference"]["top_p"],
153
+ frequency_penalty=0,
154
+ presence_penalty=0,
155
+ stop=["Question:"],
156
+ timeout=config["server"]["timeout"],
157
+ )
158
+ response_text = response.choices[0].text.strip()
159
+ else:
160
+ response = client.chat.completions.create(
161
+ model=config["server"]["model"],
162
+ messages=prompt,
163
+ temperature=config["inference"]["temperature"],
164
+ max_tokens=config["inference"]["max_tokens"],
165
+ top_p=config["inference"]["top_p"],
166
+ frequency_penalty=0,
167
+ presence_penalty=0,
168
+ stop=["Question:"],
169
+ timeout=config["server"]["timeout"],
170
+ )
171
+ response_text = response.choices[0].message.content.strip()
172
+
173
+ pred = extract_answer(response_text)
174
+ usage = response.usage
175
+
176
+ return prompt, response_text, pred, usage
177
+
178
+ except Exception as e:
179
+ logger.warning(f"Attempt {attempt + 1} failed: {e}")
180
+ if attempt == max_retries - 1:
181
+ logger.error(f"All attempts failed for question: {question['question_id']}")
182
+ return None, None, None, None
183
+ time.sleep(3) # Wait before retrying
184
+
185
+ # result_processor.py
186
+ def save_results(results, output_path, lock):
187
+ max_retries = 3
188
+ for attempt in range(max_retries):
189
+ try:
190
+ with lock:
191
+ with open(output_path, 'w') as f:
192
+ json.dump(results, f, indent=2)
193
+ return
194
+ except Exception as e:
195
+ logger.warning(f"Attempt {attempt + 1} to save results failed: {e}")
196
+ if attempt == max_retries - 1:
197
+ logger.error(f"Failed to save results to {output_path}")
198
+ time.sleep(1) # Wait before retrying
199
+
200
+ def save_summary(category_record, output_path, lock):
201
+ max_retries = 3
202
+ for attempt in range(max_retries):
203
+ try:
204
+ with lock:
205
+ with open(output_path, 'w') as f:
206
+ json.dump(category_record, f, indent=2)
207
+ return
208
+ except Exception as e:
209
+ logger.warning(f"Attempt {attempt + 1} to save summary failed: {e}")
210
+ if attempt == max_retries - 1:
211
+ logger.error(f"Failed to save summary to {output_path}")
212
+ time.sleep(1) # Wait before retrying
213
+
214
+ def update_results(results, category_record, question, pred, answer):
215
+ category = question['category']
216
+
217
+ if category not in category_record:
218
+ category_record[category] = {"correct": 0, "total": 0}
219
+
220
+ category_record[category]["total"] += 1
221
+ if pred == answer:
222
+ category_record[category]["correct"] += 1
223
+
224
+ result = {
225
+ "question_id": question['question_id'],
226
+ "category": category,
227
+ "question": question['question'],
228
+ "options": question['options'],
229
+ "pred": pred,
230
+ "answer": answer,
231
+ "correct": pred == answer
232
+ }
233
+ results.append(result)
234
+
235
+ return results, category_record
236
+
237
+ def process_and_save_results(question, pred, client, config, results, category_record, output_dir, lock):
238
+ results, category_record = update_results(results, category_record, question, pred, question['answer'])
239
+
240
+ output_res_path = os.path.join(output_dir, f"{question['category']}_result.json")
241
+ output_summary_path = os.path.join(output_dir, f"{question['category']}_summary.json")
242
+
243
+ save_results(results, output_res_path, lock)
244
+ save_summary(category_record, output_summary_path, lock)
245
+
246
+ return results, category_record
247
+
248
+ def generate_final_report(category_record, output_dir):
249
+ total_correct = sum(cat["correct"] for cat in category_record.values())
250
+ total_questions = sum(cat["total"] for cat in category_record.values())
251
+ overall_accuracy = total_correct / total_questions if total_questions > 0 else 0
252
+
253
+ report = f"MMLU-Pro Benchmark Final Report\n"
254
+ report += f"================================\n\n"
255
+ report += f"Overall Accuracy: {overall_accuracy:.2%} ({total_correct}/{total_questions})\n\n"
256
+ report += f"Category Breakdown:\n"
257
+ for category, stats in category_record.items():
258
+ accuracy = stats["correct"] / stats["total"] if stats["total"] > 0 else 0
259
+ report += f" {category}: {accuracy:.2%} ({stats['correct']}/{stats['total']})\n"
260
+
261
+ report_path = os.path.join(output_dir, "final_report.txt")
262
+ with open(report_path, 'w') as f:
263
+ f.write(report)
264
+
265
+ logger.info(f"Final report saved to {report_path}")
266
+
267
+ def mmlu_pro_main():
268
+ # Load configuration
269
+ config = load_mmlu_pro_config()
270
+
271
+ # Initialize OpenAI client
272
+ client = initialize_client(config)
273
+
274
+ # Load and preprocess the MMLU-Pro dataset
275
+ test_data, dev_data = load_mmlu_pro()
276
+ if test_data is None or dev_data is None:
277
+ logger.error("Failed to load dataset. Exiting.")
278
+ return
279
+
280
+ # Prepare output directory
281
+ output_dir = os.path.join("eval_results", config["server"]["model"].replace("/", "-"))
282
+ os.makedirs(output_dir, exist_ok=True)
283
+
284
+ # Initialize results storage
285
+ results = []
286
+ category_record = {}
287
+ lock = threading.Lock()
288
+
289
+ # Set a failure threshold to cancel the benchmark if too many questions fail
290
+ max_failed_questions = 6
291
+ failed_questions = 0
292
+
293
+ # Process each subject
294
+ for subject, questions in test_data.items():
295
+ logger.info(f"Processing subject: {subject}")
296
+ cot_examples = dev_data[subject]
297
+
298
+ # Use ThreadPoolExecutor for parallel processing
299
+ with ThreadPoolExecutor(max_workers=config["test"]["parallel"]) as executor:
300
+ futures = []
301
+ for question in questions:
302
+ future = executor.submit(run_single_question, question, cot_examples, client, config)
303
+ futures.append((future, question))
304
+
305
+ # Process results as they complete
306
+ for future, question in tqdm(futures, total=len(futures)):
307
+ prompt, response, pred, usage = future.result()
308
+
309
+ # Check if the question failed and increment the failure count
310
+ if pred is None:
311
+ failed_questions += 1
312
+ logger.warning(f"Failed question count: {failed_questions}/{max_failed_questions}")
313
+
314
+ # Stop the entire process if too many questions fail
315
+ if failed_questions >= max_failed_questions:
316
+ logger.error(f"Too many failed questions. Stopping the benchmark for {subject}.")
317
+ return
318
+
319
+ # Process and save results if the question was answered
320
+ if pred is not None:
321
+ results, category_record = process_and_save_results(
322
+ question, pred, client, config, results, category_record, output_dir, lock
323
+ )
324
+
325
+ # Save final results for the subject
326
+ save_results(results, os.path.join(output_dir, f"{subject}_final_result.json"), lock)
327
+ save_summary(category_record, os.path.join(output_dir, f"{subject}_final_summary.json"), lock)
328
+
329
+ # Generate and save final report
330
+ generate_final_report(category_record, output_dir)
331
+
332
+ logger.info(f"Evaluation complete. Results saved in {output_dir}")
333
+
334
+ def run_mmlu_pro_benchmark():
335
+ start_time = time.time()
336
+ mmlu_pro_main()
337
+ end_time = time.time()
338
+ logger.info(f"Total execution time: {end_time - start_time:.2f} seconds")
339
+ #
340
+ # End of file
341
+ ####################################################################################################
App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/config.toml ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Comment to be included in the beginning of the final report.
2
+ comment = ""
3
+
4
+ [server]
5
+ url = "http://localhost:11434/v1"
6
+ api_key = "api key"
7
+ model = "llama3"
8
+ timeout = 600.0
9
+
10
+ [inference]
11
+ # Ssettings below are from evaluate_from_local.py for VLLM on TIGER-AI-Lab/MMLU-Pro
12
+ temperature = 0.0
13
+ top_p = 1.0 # not specified but default for VLLM
14
+ max_tokens = 2048
15
+ # The variable {subject} will be replaced with appropriate value in runtime.
16
+ system_prompt = "The following are multiple choice questions (with answers) about {subject}. Think step by step and then finish your answer with \"the answer is (X)\" where X is the correct letter choice."
17
+ # "multi_chat" inserts COT examples into multi-turn messages. Use for instruct/chat models.
18
+ # "no_chat" uses v1/completion api. Use for non-instruct/chat model.
19
+ # "single_chat" (from the script for GPT-4O) inserts all the COT examples and question into a single message. Not recommended, use only for legacy compatibility.
20
+ style = "multi_chat"
21
+
22
+ [test]
23
+ categories = ['biology', 'business', 'chemistry', 'computer science', 'economics', 'engineering', 'health', 'history', 'law', 'math', 'philosophy', 'physics', 'psychology', 'other']
24
+ parallel = 1
25
+
26
+ [log]
27
+ # Verbosity between 0-2
28
+ verbosity = 0
29
+ # If true, logs exact prompt sent to the model in the test result files.
30
+ log_prompt = true
App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/mmlu_pro_test.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Test the load_config function
2
+ def test_load_config():
3
+ import sys
4
+ original_argv = sys.argv
5
+ #sys.argv = ["run_openai.py", "-c", "test_config.toml", "-u", "http://test.com", "-m", "test-model"]
6
+
7
+ config = load_config()
8
+
9
+ assert config["server"]["url"] == "http://test.com"
10
+ assert config["server"]["model"] == "test-model"
11
+
12
+ sys.argv = original_argv
13
+ print("load_config test passed")
14
+
15
+ def test_load_mmlu_pro():
16
+ test_df, val_df = load_mmlu_pro()
17
+ assert test_df is not None
18
+ assert val_df is not None
19
+ assert isinstance(test_df, dict)
20
+ assert isinstance(val_df, dict)
21
+ print("load_mmlu_pro test passed")
22
+
23
+
24
+ def test_initialize_client():
25
+ test_config = {
26
+ "server": {
27
+ "url": "http://test.com",
28
+ "api_key": "test_key",
29
+ "timeout": 30
30
+ }
31
+ }
32
+
33
+ client = initialize_client(test_config)
34
+
35
+ assert client.base_url == "http://test.com"
36
+ assert client.api_key == "test_key"
37
+ assert client.timeout == 30
38
+
39
+ print("initialize_client test passed")
40
+
41
+
42
+ test_initialize_client()
43
+
44
+ def test_preprocess():
45
+ sample_data = [
46
+ {"category": "math", "options": ["A", "B", "N/A", "C"]},
47
+ {"category": "science", "options": ["X", "Y", "Z"]}
48
+ ]
49
+ processed = preprocess(sample_data)
50
+ assert "math" in processed
51
+ assert "science" in processed
52
+ assert len(processed["math"][0]["options"]) == 3
53
+ assert "N/A" not in processed["math"][0]["options"]
54
+ assert len(processed["science"][0]["options"]) == 3
55
+ print("preprocess test passed")
56
+
57
+ test_load_mmlu_pro()
58
+ test_preprocess()
59
+
60
+
61
+ test_load_config()
62
+
63
+
64
+ def test_create_prompt():
65
+ config = {
66
+ "inference": {
67
+ "style": "multi_chat",
68
+ "system_prompt": "You are a helpful assistant."
69
+ }
70
+ }
71
+ cot_examples = [{
72
+ "question": "What is 2+2?",
73
+ "options": ["3", "4", "5"],
74
+ "cot_content": "Let's add 2 and 2. 2+2 = 4."
75
+ }]
76
+ question = "What is 3+3?"
77
+ options = ["5", "6", "7"]
78
+
79
+ # Test multi_chat
80
+ result = create_prompt(cot_examples, question, options, config)
81
+ assert isinstance(result, list)
82
+ assert len(result) == 4
83
+ assert result[0]["role"] == "system"
84
+ assert result[-1]["role"] == "user"
85
+
86
+ # Test single_chat
87
+ config["inference"]["style"] = "single_chat"
88
+ result = create_prompt(cot_examples, question, options, config)
89
+ assert isinstance(result, list)
90
+ assert len(result) == 1
91
+ assert result[0]["role"] == "user"
92
+
93
+ # Test no_chat
94
+ config["inference"]["style"] = "no_chat"
95
+ result = create_prompt(cot_examples, question, options, config)
96
+ assert isinstance(result, str)
97
+ assert "What is 2+2?" in result
98
+ assert "What is 3+3?" in result
99
+
100
+ print("create_prompt test passed")
101
+
102
+ test_create_prompt()
103
+
104
+
105
+ def test_extract_answer():
106
+ test_cases = [
107
+ ("The answer is (B)", "B"),
108
+ ("After careful consideration, I believe the answer is C.", "C"),
109
+ (
110
+ "Let's analyze each option:\nA. Incorrect\nB. Incorrect\nC. Correct\nD. Incorrect\nTherefore, the answer is C.",
111
+ "C"),
112
+ ("A. GHTIS\nB. MCU\nC. UBT\nD. ALIN\n\nThe correct answer is B. MCU.", "B"),
113
+ ("There is no clear answer in this text.", None),
114
+ ("The options are A, B, C, and D. I think B is the best answer.", "B")
115
+ ]
116
+
117
+ for text, expected in test_cases:
118
+ result = extract_answer(text)
119
+ assert result == expected, f"Failed on input '{text}'. Expected {expected}, got {result}"
120
+
121
+ print("extract_answer test passed")
122
+
123
+
124
+ test_extract_answer()
125
+
126
+ from unittest.mock import Mock
127
+
128
+ def test_run_single_question():
129
+ # Mock OpenAI client
130
+ mock_client = Mock()
131
+ mock_response = Mock()
132
+ mock_response.choices = [Mock(text="The answer is B", message=Mock(content="The answer is B"))]
133
+ mock_response.usage = Mock(prompt_tokens=10, completion_tokens=20, total_tokens=30)
134
+ mock_client.completions.create.return_value = mock_response
135
+ mock_client.chat.completions.create.return_value = mock_response
136
+
137
+ # Mock configuration
138
+ config = {
139
+ "inference": {
140
+ "style": "no_chat",
141
+ "system_prompt": "You are a helpful assistant.",
142
+ "temperature": 0.7,
143
+ "max_tokens": 100,
144
+ "top_p": 1.0
145
+ },
146
+ "server": {
147
+ "model": "test-model",
148
+ "timeout": 30
149
+ }
150
+ }
151
+
152
+ # Mock question and examples
153
+ question = {
154
+ "question": "What is 2+2?",
155
+ "options": ["3", "4", "5"]
156
+ }
157
+ cot_examples = []
158
+
159
+ # Test no_chat style
160
+ prompt, response, pred, usage = run_single_question(question, cot_examples, mock_client, config)
161
+ assert prompt is not None
162
+ assert response == "The answer is B"
163
+ assert pred == "B"
164
+ assert usage.prompt_tokens == 10
165
+ assert usage.completion_tokens == 20
166
+ assert usage.total_tokens == 30
167
+
168
+ # Test chat style
169
+ config["inference"]["style"] = "multi_chat"
170
+ prompt, response, pred, usage = run_single_question(question, cot_examples, mock_client, config)
171
+ assert prompt is not None
172
+ assert response == "The answer is B"
173
+ assert pred == "B"
174
+ assert usage.prompt_tokens == 10
175
+ assert usage.completion_tokens == 20
176
+ assert usage.total_tokens == 30
177
+
178
+ print("run_single_question test passed")
179
+
180
+ test_run_single_question()
181
+
182
+
183
+ def test_save_and_update_functions():
184
+ # Create a temporary directory for test files
185
+ with tempfile.TemporaryDirectory() as tmpdir:
186
+ lock = threading.Lock()
187
+ results = []
188
+ category_record = {}
189
+
190
+ # Test question
191
+ question = {
192
+ 'question_id': '1',
193
+ 'category': 'math',
194
+ 'question': 'What is 2+2?',
195
+ 'options': ['3', '4', '5'],
196
+ 'answer': 'B'
197
+ }
198
+
199
+ # Test update_results
200
+ results, category_record = update_results(results, category_record, question, 'B', 'B')
201
+ assert len(results) == 1
202
+ assert category_record['math']['correct'] == 1
203
+ assert category_record['math']['total'] == 1
204
+
205
+ # Test save_results and save_summary
206
+ results_path = os.path.join(tmpdir, 'results.json')
207
+ summary_path = os.path.join(tmpdir, 'summary.json')
208
+
209
+ save_results(results, results_path, lock)
210
+ save_summary(category_record, summary_path, lock)
211
+
212
+ assert os.path.exists(results_path)
213
+ assert os.path.exists(summary_path)
214
+
215
+ # Test process_and_save_results
216
+ config = {'server': {'model': 'test-model'}}
217
+ client = None # We don't need a real client for this test
218
+
219
+ results, category_record = process_and_save_results(question, 'B', client, config, results, category_record,
220
+ tmpdir, lock)
221
+
222
+ assert len(results) == 2
223
+ assert category_record['math']['correct'] == 2
224
+ assert category_record['math']['total'] == 2
225
+
226
+ assert os.path.exists(os.path.join(tmpdir, 'math_result.json'))
227
+ assert os.path.exists(os.path.join(tmpdir, 'math_summary.json'))
228
+
229
+ print("save_and_update_functions tests passed")
230
+
231
+
232
+ test_save_and_update_functions()
App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/run_openai.py ADDED
@@ -0,0 +1,546 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Script taken from: https://github.com/chigkim/Ollama-MMLU-Pro
2
+ # No changes made
3
+ import os
4
+ import re
5
+ import json
6
+ import time
7
+ import random
8
+ from tqdm import tqdm
9
+ from openai import OpenAI
10
+ from datasets import load_dataset
11
+ from concurrent.futures import ThreadPoolExecutor, as_completed
12
+ import threading
13
+ from datetime import datetime, timedelta
14
+ import codecs
15
+ import toml
16
+ import argparse
17
+ import queue
18
+ import numpy as np
19
+ import copy
20
+
21
+ parser = argparse.ArgumentParser(
22
+ prog="python3 run_openai.py",
23
+ description="Run MMLU Pro Benchmark for a local LLM via OpenAI Compatible API.",
24
+ epilog="Specify options above to override one or more settings from config.",
25
+ )
26
+ parser.add_argument(
27
+ "-c",
28
+ "--config",
29
+ help="Configuration file. Default=config.toml",
30
+ default="config.toml",
31
+ )
32
+ parser.add_argument(
33
+ "-u",
34
+ "--url",
35
+ help="server url",
36
+ )
37
+ parser.add_argument("-a", "--api", help="api key")
38
+ parser.add_argument("-m", "--model", help="Model name")
39
+ parser.add_argument(
40
+ "--timeout",
41
+ type=float,
42
+ help="Request timeout in seconds",
43
+ )
44
+ parser.add_argument("--category", type=str)
45
+ parser.add_argument("-p", "--parallel", type=int, help="Number of parallel requests")
46
+ parser.add_argument("-v", "--verbosity", type=int, help="Verbosity level 0-2")
47
+ parser.add_argument(
48
+ "--log_prompt",
49
+ help="Writes exact prompt and response into log.txt",
50
+ action="store_true",
51
+ )
52
+ parser.add_argument(
53
+ "--comment", type=str, help="Comment to be included in the final report."
54
+ )
55
+ args = parser.parse_args()
56
+ config = toml.load(open(args.config))
57
+ if args.url:
58
+ config["server"]["url"] = args.url
59
+ if args.api:
60
+ config["server"]["api_key"] = args.api
61
+ if args.model:
62
+ config["server"]["model"] = args.model
63
+ if args.timeout:
64
+ config["server"]["timeout"] = args.timeout
65
+ if args.category:
66
+ config["test"]["categories"] = [args.category]
67
+ if args.parallel:
68
+ config["test"]["parallel"] = args.parallel
69
+ if args.verbosity:
70
+ config["log"]["verbosity"] = args.verbosity
71
+ if args.log_prompt:
72
+ config["log"]["log_prompt"] = args.log_prompt
73
+ if args.comment:
74
+ config["comment"] = args.comment
75
+
76
+
77
+ client = OpenAI(
78
+ base_url=config["server"]["url"],
79
+ api_key=config["server"]["api_key"],
80
+ timeout=config["server"]["timeout"],
81
+ )
82
+
83
+
84
+ def log(message):
85
+ print(message)
86
+ with codecs.open(log_path, "a", "utf-8") as file:
87
+ file.write(message + "\n")
88
+
89
+
90
+ def get_chat_completion(messages):
91
+ try:
92
+ response = client.chat.completions.create(
93
+ model=config["server"]["model"],
94
+ messages=messages,
95
+ temperature=config["inference"]["temperature"],
96
+ max_tokens=config["inference"]["max_tokens"],
97
+ top_p=config["inference"]["top_p"],
98
+ frequency_penalty=0,
99
+ presence_penalty=0,
100
+ stop=["Question:"],
101
+ timeout=config["server"]["timeout"],
102
+ )
103
+ try:
104
+ usage_q.put(
105
+ (response.usage.prompt_tokens, response.usage.completion_tokens)
106
+ )
107
+ except:
108
+ pass
109
+ return response.choices[0].message.content.strip()
110
+ except Exception as e:
111
+ print("Resubmitting, Error: ", e)
112
+ time.sleep(3)
113
+ return get_chat_completion(messages)
114
+
115
+
116
+ def get_completion(prompt):
117
+ try:
118
+ response = client.completions.create(
119
+ model=config["server"]["model"],
120
+ prompt=prompt,
121
+ temperature=config["inference"]["temperature"],
122
+ max_tokens=config["inference"]["max_tokens"],
123
+ top_p=config["inference"]["top_p"],
124
+ frequency_penalty=0,
125
+ presence_penalty=0,
126
+ stop=["Question:"],
127
+ timeout=config["server"]["timeout"],
128
+ )
129
+ try:
130
+ usage_q.put(
131
+ (response.usage.prompt_tokens, response.usage.completion_tokens)
132
+ )
133
+ except:
134
+ pass
135
+ if response.choices:
136
+ return response.choices[0].text.strip()
137
+ elif response.content:
138
+ return response.content.strip()
139
+ print("Can't get response.")
140
+ return None
141
+ except Exception as e:
142
+ print("Resubmitting, Error: ", e)
143
+ time.sleep(3)
144
+ return get_completion(prompt)
145
+
146
+
147
+ def load_mmlu_pro():
148
+ dataset = load_dataset("TIGER-Lab/MMLU-Pro")
149
+ test_df, val_df = dataset["test"], dataset["validation"]
150
+ test_df = preprocess(test_df)
151
+ val_df = preprocess(val_df)
152
+ return test_df, val_df
153
+
154
+
155
+ def preprocess(test_df):
156
+ res_df = []
157
+ for each in test_df:
158
+ options = []
159
+ for opt in each["options"]:
160
+ if opt == "N/A":
161
+ continue
162
+ options.append(opt)
163
+ each["options"] = options
164
+ res_df.append(each)
165
+ res = {}
166
+ for each in res_df:
167
+ if each["category"] not in res:
168
+ res[each["category"]] = []
169
+ res[each["category"]].append(each)
170
+ return res
171
+
172
+
173
+ def format_example(question, options, cot_content=""):
174
+ if cot_content == "":
175
+ cot_content = "Let's think step by step."
176
+ if cot_content.startswith("A: "):
177
+ cot_content = cot_content[3:]
178
+ example = "Question: {}\nOptions: ".format(question)
179
+ choice_map = "ABCDEFGHIJ"
180
+ for i, opt in enumerate(options):
181
+ example += "{}. {}\n".format(choice_map[i], opt)
182
+ return example.strip(), cot_content.strip()
183
+
184
+
185
+ def multi_chat_prompt(cot_examples, question, options):
186
+ messages = [
187
+ {
188
+ "role": "system",
189
+ "content": config["inference"]["system_prompt"],
190
+ },
191
+ ]
192
+ for each in cot_examples:
193
+ example, cot_content = format_example(
194
+ each["question"], each["options"], each["cot_content"]
195
+ )
196
+ messages.append({"role": "user", "content": example})
197
+ messages.append({"role": "assistant", "content": "Answer: " + cot_content})
198
+ example, cot_content = format_example(question, options)
199
+ messages.append({"role": "user", "content": example})
200
+ return messages
201
+
202
+
203
+ def single_chat_prompt(cot_examples, question, options):
204
+ messages = [
205
+ {
206
+ "role": "system",
207
+ "content": config["inference"]["system_prompt"],
208
+ },
209
+ ]
210
+ prompt = no_chat_prompt(cot_examples, question, options, no_system=True)
211
+ messages.append({"role": "user", "content": prompt})
212
+ return messages
213
+
214
+
215
+ def no_chat_prompt(cot_examples, question, options, no_system=False):
216
+ prompt = config["inference"]["system_prompt"] + "\n\n"
217
+ if no_system:
218
+ prompt = ""
219
+ for each in cot_examples:
220
+ example, cot_content = format_example(
221
+ each["question"], each["options"], each["cot_content"]
222
+ )
223
+ prompt += example + "\n"
224
+ prompt += "Answer: " + cot_content + "\n\n"
225
+ example, cot_content = format_example(question, options)
226
+ prompt += example + "\n"
227
+ prompt += "Answer: " + cot_content
228
+ return prompt
229
+
230
+
231
+ def extract_answer(text):
232
+ pattern = r"answer is \(?([ABCDEFGHIJ])\)?"
233
+ match = re.search(pattern, text)
234
+ if match:
235
+ return match.group(1)
236
+ else:
237
+ return extract_again(text)
238
+
239
+
240
+ def extract_again(text):
241
+ pattern = r".*[aA]nswer:\s*\(?([A-J])\)?"
242
+ match = re.search(pattern, text)
243
+ if match:
244
+ return match.group(1)
245
+ else:
246
+ return extract_final(text)
247
+
248
+
249
+ def extract_final(text):
250
+ pattern = r"\b[A-J]\b(?!.*\b[A-J]\b)"
251
+ match = re.search(pattern, text, re.DOTALL)
252
+ if match:
253
+ return match[0]
254
+ else:
255
+ if config["log"]["verbosity"] >= 1:
256
+ print("Extraction failed:\n", text)
257
+ return None
258
+
259
+
260
+ def run_single_question(single_question, cot_examples_dict, exist_result):
261
+ exist = True
262
+ q_id = single_question["question_id"]
263
+ for each in exist_result:
264
+ if (
265
+ q_id == each["question_id"]
266
+ and single_question["question"] == each["question"]
267
+ ):
268
+ if config["log"]["verbosity"] >= 1:
269
+ print("already exists, skipping.")
270
+ return None, None, None, exist
271
+ exist = False
272
+ category = single_question["category"]
273
+ cot_examples = cot_examples_dict[category]
274
+ question = single_question["question"]
275
+ options = single_question["options"]
276
+ try:
277
+ if config["inference"]["style"] == "single_chat":
278
+ prompt = single_chat_prompt(cot_examples, question, options)
279
+ response = get_chat_completion(prompt)
280
+ elif config["inference"]["style"] == "multi_chat":
281
+ prompt = multi_chat_prompt(cot_examples, question, options)
282
+ response = get_chat_completion(prompt)
283
+ elif config["inference"]["style"] == "no_chat":
284
+ prompt = no_chat_prompt(cot_examples, question, options)
285
+ response = get_completion(prompt)
286
+ except Exception as e:
287
+ print("error", e)
288
+ return None, None, None, exist
289
+ pred = extract_answer(response)
290
+ return prompt, response, pred, exist
291
+
292
+
293
+ def update_result(output_res_path, lock):
294
+ category_record = {}
295
+ res = []
296
+ success = False
297
+ while not success:
298
+ try:
299
+ if os.path.exists(output_res_path):
300
+ with lock:
301
+ with open(output_res_path, "r") as fi:
302
+ res = json.load(fi)
303
+ for each in res:
304
+ category = each["category"]
305
+ if category not in category_record:
306
+ category_record[category] = {"corr": 0.0, "wrong": 0.0}
307
+ category_record["random"] = {"corr": 0.0, "wrong": 0.0}
308
+ if not each["pred"]:
309
+ random.seed(12345)
310
+ x = random.randint(0, len(each["options"]) - 1)
311
+ if x == each["answer_index"]:
312
+ category_record[category]["corr"] += 1
313
+ category_record["random"]["corr"] += 1
314
+ else:
315
+ category_record[category]["wrong"] += 1
316
+ category_record["random"]["wrong"] += 1
317
+ elif each["pred"] == each["answer"]:
318
+ category_record[category]["corr"] += 1
319
+ else:
320
+ category_record[category]["wrong"] += 1
321
+ success = True
322
+ except Exception as e:
323
+ print("Error", e)
324
+ return res, category_record
325
+
326
+
327
+ def evaluate(subjects):
328
+ test_df, dev_df = load_mmlu_pro()
329
+ if not subjects:
330
+ subjects = list(test_df.keys())
331
+ print("assigned subjects", subjects)
332
+ lock = threading.Lock()
333
+ system_prompt = config["inference"]["system_prompt"]
334
+ for subject in subjects:
335
+ start = time.time()
336
+ print(f"Testing {subject}...")
337
+ config["inference"]["system_prompt"] = system_prompt.replace(
338
+ "{subject}", subject
339
+ )
340
+ test_data = test_df[subject]
341
+ output_res_path = os.path.join(output_dir, subject + "_result.json")
342
+ output_summary_path = os.path.join(output_dir, subject + "_summary.json")
343
+ res, category_record = update_result(output_res_path, lock)
344
+
345
+ with ThreadPoolExecutor(max_workers=config["test"]["parallel"]) as executor:
346
+ futures = {
347
+ executor.submit(run_single_question, each, dev_df, res): each
348
+ for each in test_data
349
+ }
350
+ for future in tqdm(
351
+ as_completed(futures), total=len(futures), smoothing=0.0, ascii=True
352
+ ):
353
+ each = futures[future]
354
+ label = each["answer"]
355
+ category = subject
356
+ prompt, response, pred, exist = future.result()
357
+ if exist:
358
+ continue
359
+ if response is not None:
360
+ res, category_record = update_result(output_res_path, lock)
361
+ if category not in category_record:
362
+ category_record[category] = {"corr": 0.0, "wrong": 0.0}
363
+ if config["log"]["log_prompt"]:
364
+ each["prompt"] = prompt
365
+ each["response"] = response
366
+ each["pred"] = pred
367
+ res.append(each)
368
+ if config["log"]["verbosity"] >= 2:
369
+ log_json = {
370
+ "id": each["question_id"],
371
+ "question": each["question"],
372
+ "response": each["response"],
373
+ "pred": each["pred"],
374
+ "answer": each["answer"],
375
+ }
376
+ print("\n" + json.dumps(log_json, indent="\t"))
377
+ if pred is not None:
378
+ if pred == label:
379
+ category_record[category]["corr"] += 1
380
+ else:
381
+ category_record[category]["wrong"] += 1
382
+ else:
383
+ category_record[category]["wrong"] += 1
384
+ save_res(res, output_res_path, lock)
385
+ save_summary(category_record, output_summary_path, lock)
386
+ res, category_record = update_result(output_res_path, lock)
387
+ save_res(res, output_res_path, lock)
388
+ hours, minutes, seconds = elapsed(start)
389
+ log(
390
+ f"Finished testing {subject} in {hours} hours, {minutes} minutes, {seconds} seconds."
391
+ )
392
+ save_summary(category_record, output_summary_path, lock, report=True)
393
+
394
+
395
+ def save_res(res, output_res_path, lock):
396
+ temp = []
397
+ exist_q_id = []
398
+ for each in res:
399
+ if each["question_id"] not in exist_q_id:
400
+ exist_q_id.append(each["question_id"])
401
+ temp.append(each)
402
+ else:
403
+ continue
404
+ res = temp
405
+ with lock:
406
+ with open(output_res_path, "w") as fo:
407
+ fo.write(json.dumps(res, indent="\t"))
408
+
409
+
410
+ def print_score(label, corr, wrong):
411
+ try:
412
+ corr = int(corr)
413
+ wrong = int(wrong)
414
+ total = corr + wrong
415
+ acc = corr / total * 100
416
+ log(f"{label}, {corr}/{total}, {acc:.2f}%")
417
+ except Exception as e:
418
+ log(f"{label}, {e} error")
419
+
420
+
421
+ def save_summary(category_record, output_summary_path, lock, report=False):
422
+ total_corr = 0.0
423
+ total_wrong = 0.0
424
+ for k, v in category_record.items():
425
+ if k == "total" or k == "random":
426
+ continue
427
+ cat_acc = v["corr"] / (v["corr"] + v["wrong"])
428
+ category_record[k]["acc"] = cat_acc
429
+ total_corr += v["corr"]
430
+ total_wrong += v["wrong"]
431
+ acc = total_corr / (total_corr + total_wrong)
432
+ category_record["total"] = {"corr": total_corr, "wrong": total_wrong, "acc": acc}
433
+ if report:
434
+ print_score("Total", total_corr, total_wrong)
435
+ if "random" in category_record:
436
+ random_corr = category_record["random"]["corr"]
437
+ random_wrong = category_record["random"]["wrong"]
438
+ print_score(
439
+ "Random Guess Attempts",
440
+ random_corr + random_wrong,
441
+ total_corr + total_wrong - random_corr - random_wrong,
442
+ )
443
+ print_score("Correct Random Guesses", random_corr, random_wrong)
444
+ print_score(
445
+ "Adjusted Score Without Random Guesses",
446
+ total_corr - random_corr,
447
+ total_wrong - random_wrong,
448
+ )
449
+ with lock:
450
+ with open(output_summary_path, "w") as fo:
451
+ fo.write(json.dumps(category_record, indent="\t"))
452
+
453
+
454
+ def final_report(assigned_subjects):
455
+ total_corr = 0.0
456
+ total_wrong = 0.0
457
+ random_corr = 0.0
458
+ random_wrong = 0.0
459
+ names = ["overall"] + assigned_subjects
460
+ table = "| " + " | ".join(names) + " |\n"
461
+ separators = [re.sub(r".", "-", name) for name in names]
462
+ table += "| " + " | ".join(separators) + " |\n"
463
+ scores = []
464
+ for file in assigned_subjects:
465
+ res = json.load(open(os.path.join(output_dir, file + "_summary.json")))
466
+ cat_corr = res["total"]["corr"]
467
+ total_corr += cat_corr
468
+ cat_wrong = res["total"]["wrong"]
469
+ total_wrong += cat_wrong
470
+ scores.append(cat_corr / (cat_corr + cat_wrong))
471
+ if "random" in res:
472
+ random_corr += res["random"]["corr"]
473
+ random_wrong += res["random"]["wrong"]
474
+ print_score("Total", total_corr, total_wrong)
475
+ if random_corr and random_wrong:
476
+ print_score(
477
+ "Random Guess Attempts",
478
+ random_corr + random_wrong,
479
+ total_corr + total_wrong - random_corr - random_wrong,
480
+ )
481
+ print_score("Correct Random Guesses", random_corr, random_wrong)
482
+ print_score(
483
+ "Adjusted Score Without Random Guesses",
484
+ total_corr - random_corr,
485
+ total_wrong - random_wrong,
486
+ )
487
+ scores.insert(0, total_corr / (total_corr + total_wrong))
488
+ scores = [f"{score*100:.2f}" for score in scores]
489
+ table += "| " + " | ".join(scores) + " |"
490
+ token_report()
491
+ log("Markdown Table:")
492
+ log(table)
493
+
494
+
495
+ def elapsed(start):
496
+ duration = time.time() - start
497
+ duration_td = timedelta(seconds=duration)
498
+ hours, remainder = divmod(duration_td.seconds, 3600)
499
+ minutes, seconds = divmod(remainder, 60)
500
+ return hours, minutes, seconds
501
+
502
+
503
+ def token_report():
504
+ ptoks = []
505
+ ctoks = []
506
+ while not usage_q.empty():
507
+ usage = usage_q.get()
508
+ ptoks.append(usage[0])
509
+ ctoks.append(usage[1])
510
+ if ptoks and ctoks:
511
+ log("Token Usage:")
512
+ duration = end - start
513
+ ptoks = np.array(ptoks)
514
+ ctoks = np.array(ctoks)
515
+ log(
516
+ f"Prompt tokens: min {ptoks.min()}, average {ptoks.mean():.0f}, max {ptoks.max()}, total {ptoks.sum()}, tk/s {ptoks.sum()/duration:.2f}"
517
+ )
518
+ log(
519
+ f"Completion tokens: min {ctoks.min()}, average {ctoks.mean():.0f}, max {ctoks.max()}, total {ctoks.sum()}, tk/s {ctoks.sum()/duration:.2f}"
520
+ )
521
+
522
+
523
+ if __name__ == "__main__":
524
+ usage_q = queue.Queue()
525
+ output_dir = "eval_results/" + re.sub(r"\W", "-", config["server"]["model"])
526
+ os.makedirs(output_dir, exist_ok=True)
527
+ log_path = os.path.join(output_dir, "report.txt")
528
+ try:
529
+ os.remove(log_path)
530
+ except:
531
+ pass
532
+ config_copy = copy.deepcopy(config)
533
+ del config_copy["server"]["api_key"]
534
+ del config_copy["test"]["categories"]
535
+ log(f"{datetime.now()}")
536
+ log(json.dumps(config_copy, indent="\t"))
537
+ assigned_subjects = config["test"]["categories"]
538
+ start = time.time()
539
+ evaluate(assigned_subjects)
540
+ end = time.time()
541
+ hours, minutes, seconds = elapsed(start)
542
+ log(
543
+ f"Finished the benchmark in {hours} hours, {minutes} minutes, {seconds} seconds."
544
+ )
545
+ final_report(assigned_subjects)
546
+ print("Report saved to:", log_path)
App_Function_Libraries/Benchmarks_Evaluations/__init__.py ADDED
File without changes
App_Function_Libraries/Benchmarks_Evaluations/ms_g_eval.py ADDED
@@ -0,0 +1,498 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #######################################################################################################################
2
+ #
3
+ # Evaluations_Benchmarks_tab.py
4
+ #
5
+ # Description: This file contains the code to evaluate the generated text using G-Eval metric.
6
+ #
7
+ # Scripts taken from https://github.com/microsoft/promptflow/tree/main/examples/flows/evaluation/eval-summarization and modified.
8
+ #
9
+ import configparser
10
+ import inspect
11
+ import json
12
+ import logging
13
+ import os
14
+ import re
15
+ from typing import Dict, Callable, List, Any
16
+
17
+ import gradio as gr
18
+ from tenacity import (
19
+ RetryError,
20
+ Retrying,
21
+ after_log,
22
+ before_sleep_log,
23
+ stop_after_attempt,
24
+ wait_random_exponential,
25
+ )
26
+
27
+ from App_Function_Libraries.Chat import chat_api_call
28
+
29
+ #
30
+ #######################################################################################################################
31
+ #
32
+ # Start of G-Eval.py
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+ current_dir = os.path.dirname(os.path.abspath(__file__))
37
+ # Construct the path to the config file
38
+ config_path = os.path.join(current_dir, 'Config_Files', 'config.txt')
39
+ # Read the config file
40
+ config = configparser.ConfigParser()
41
+ config.read(config_path)
42
+
43
+
44
+ def aggregate(
45
+ fluency_list: List[float],
46
+ consistency_list: List[float],
47
+ relevance_list: List[float],
48
+ coherence_list: List[float],
49
+ ) -> Dict[str, float]:
50
+ """
51
+ Takes list of scores for 4 dims and outputs average for them.
52
+
53
+ Args:
54
+ fluency_list (List(float)): list of fluency scores
55
+ consistency_list (List(float)): list of consistency scores
56
+ relevance_list (List(float)): list of relevance scores
57
+ coherence_list (List(float)): list of coherence scores
58
+
59
+ Returns:
60
+ Dict[str, float]: Returns average scores
61
+ """
62
+ average_fluency = sum(fluency_list) / len(fluency_list)
63
+ average_consistency = sum(consistency_list) / len(consistency_list)
64
+ average_relevance = sum(relevance_list) / len(relevance_list)
65
+ average_coherence = sum(coherence_list) / len(coherence_list)
66
+
67
+ log_metric("average_fluency", average_fluency)
68
+ log_metric("average_consistency", average_consistency)
69
+ log_metric("average_relevance", average_relevance)
70
+ log_metric("average_coherence", average_coherence)
71
+
72
+ return {
73
+ "average_fluency": average_fluency,
74
+ "average_consistency": average_consistency,
75
+ "average_relevance": average_relevance,
76
+ "average_coherence": average_coherence,
77
+ }
78
+
79
+ def run_geval(transcript: str, summary: str, api_key: str, api_name: str = None, save: bool = False):
80
+ try:
81
+ validate_inputs(transcript, summary, api_name, api_key)
82
+ except ValueError as e:
83
+ return str(e)
84
+
85
+ prompts = {
86
+ "coherence": """You will be given one summary written for a source document.
87
+
88
+ Your task is to rate the summary on one metric.
89
+
90
+ Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.
91
+
92
+ Evaluation Criteria:
93
+
94
+ Coherence (1-5) - the collective quality of all sentences. We align this dimension with the DUC quality question of structure and coherence whereby "the summary should be well-structured and well-organized. The summary should not just be a heap of related information, but should build from sentence to a coherent body of information about a topic."
95
+
96
+ Evaluation Steps:
97
+
98
+ 1. Read the source document carefully and identify the main topic and key points.
99
+ 2. Read the summary and compare it to the source document. Check if the summary covers the main topic and key points of the source document, and if it presents them in a clear and logical order.
100
+ 3. Assign a score for coherence on a scale of 1 to 5, where 1 is the lowest and 5 is the highest based on the Evaluation Criteria.
101
+
102
+
103
+ Example:
104
+
105
+
106
+ Source Document:
107
+
108
+ {{Document}}
109
+
110
+ Summary:
111
+
112
+ {{Summary}}
113
+
114
+
115
+ Evaluation Form (scores ONLY):
116
+
117
+ - Coherence:""",
118
+ "consistency": """You will be given a source document. You will then be given one summary written for this source document.
119
+
120
+ Your task is to rate the summary on one metric.
121
+
122
+ Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.
123
+
124
+
125
+ Evaluation Criteria:
126
+
127
+ Consistency (1-5) - the factual alignment between the summary and the summarized source. A factually consistent summary contains only statements that are entailed by the source document. Annotators were also asked to penalize summaries that contained hallucinated facts.
128
+
129
+ Evaluation Steps:
130
+
131
+ 1. Read the source document carefully and identify the main facts and details it presents.
132
+ 2. Read the summary and compare it to the source document. Check if the summary contains any factual errors that are not supported by the source document.
133
+ 3. Assign a score for consistency based on the Evaluation Criteria.
134
+
135
+
136
+ Example:
137
+
138
+
139
+ Source Document:
140
+
141
+ {{Document}}
142
+
143
+ Summary:
144
+
145
+ {{Summary}}
146
+
147
+
148
+ Evaluation Form (scores ONLY):
149
+
150
+ - Consistency:""",
151
+ "fluency": """You will be given one summary written for a source document.
152
+
153
+ Your task is to rate the summary on one metric.
154
+
155
+ Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.
156
+
157
+
158
+ Evaluation Criteria:
159
+
160
+ Fluency (1-3): the quality of the summary in terms of grammar, spelling, punctuation, word choice, and sentence structure.
161
+
162
+ - 1: Poor. The summary has many errors that make it hard to understand or sound unnatural.
163
+ - 2: Fair. The summary has some errors that affect the clarity or smoothness of the text, but the main points are still comprehensible.
164
+ - 3: Good. The summary has few or no errors and is easy to read and follow.
165
+
166
+
167
+ Example:
168
+
169
+ Summary:
170
+
171
+ {{Summary}}
172
+
173
+
174
+ Evaluation Form (scores ONLY):
175
+
176
+ - Fluency (1-3):""",
177
+ "relevance": """You will be given one summary written for a source document.
178
+
179
+ Your task is to rate the summary on one metric.
180
+
181
+ Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.
182
+
183
+ Evaluation Criteria:
184
+
185
+ Relevance (1-5) - selection of important content from the source. The summary should include only important information from the source document. Annotators were instructed to penalize summaries which contained redundancies and excess information.
186
+
187
+ Evaluation Steps:
188
+
189
+ 1. Read the summary and the source document carefully.
190
+ 2. Compare the summary to the source document and identify the main points of the source document.
191
+ 3. Assess how well the summary covers the main points of the source document, and how much irrelevant or redundant information it contains.
192
+ 4. Assign a relevance score from 1 to 5.
193
+
194
+
195
+ Example:
196
+
197
+
198
+ Source Document:
199
+
200
+ {{Document}}
201
+
202
+ Summary:
203
+
204
+ {{Summary}}
205
+
206
+
207
+ Evaluation Form (scores ONLY):
208
+
209
+ - Relevance:"""
210
+ }
211
+
212
+ scores = {}
213
+ explanations = {}
214
+ for metric, prompt in prompts.items():
215
+ full_prompt = prompt.replace("{{Document}}", transcript).replace("{{Summary}}", summary)
216
+ try:
217
+ score = geval_summarization(full_prompt, 5 if metric != "fluency" else 3, api_name, api_key)
218
+ scores[metric] = score
219
+ explanations[metric] = "Score based on the evaluation criteria."
220
+ except Exception as e:
221
+ error_message = detailed_api_error(api_name, e)
222
+ return error_message
223
+
224
+ avg_scores = aggregate([scores['fluency']], [scores['consistency']],
225
+ [scores['relevance']], [scores['coherence']])
226
+
227
+ results = {
228
+ "scores": scores,
229
+ "average_scores": avg_scores
230
+ }
231
+ logging.debug("Results: %s", results)
232
+
233
+ if save is not None:
234
+ logging.debug("Saving results to geval_results.json")
235
+ save_eval_results(results)
236
+ logging.debug("Results saved to geval_results.json")
237
+
238
+ formatted_result = f"""
239
+ Confabulation Check Results:
240
+
241
+ Coherence: {scores['coherence']:.2f} - {explanations['coherence']}
242
+ Consistency: {scores['consistency']:.2f} - {explanations['consistency']}
243
+ Fluency: {scores['fluency']:.2f} - {explanations['fluency']}
244
+ Relevance: {scores['relevance']:.2f} - {explanations['relevance']}
245
+
246
+ Overall Assessment: The summary has been evaluated on four key metrics.
247
+ The average scores are:
248
+ Fluency: {avg_scores['average_fluency']:.2f}
249
+ Consistency: {avg_scores['average_consistency']:.2f}
250
+ Relevance: {avg_scores['average_relevance']:.2f}
251
+ Coherence: {avg_scores['average_coherence']:.2f}
252
+
253
+ These scores indicate the overall quality of the summary in terms of its
254
+ coherence, consistency with the original text, fluency of language, and
255
+ relevance of content.
256
+ """
257
+
258
+ return formatted_result
259
+
260
+
261
+ def create_geval_tab():
262
+ with gr.Tab("G-Eval", id="g-eval"):
263
+ gr.Markdown("# G-Eval Summarization Evaluation")
264
+ with gr.Row():
265
+ with gr.Column():
266
+ document_input = gr.Textbox(label="Source Document", lines=10)
267
+ summary_input = gr.Textbox(label="Summary", lines=5)
268
+ api_name_input = gr.Dropdown(
269
+ choices=["OpenAI", "Anthropic", "Cohere", "Groq", "OpenRouter", "DeepSeek", "HuggingFace", "Mistral", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "Local-LLM", "Ollama"],
270
+ label="Select API"
271
+ )
272
+ api_key_input = gr.Textbox(label="API Key (if required)", type="password")
273
+ save_value = gr.Checkbox(label="Save Results to a JSON file(geval_results.json)")
274
+ evaluate_button = gr.Button("Evaluate Summary")
275
+ with gr.Column():
276
+ output = gr.Textbox(label="Evaluation Results", lines=10)
277
+
278
+ evaluate_button.click(
279
+ fn=run_geval,
280
+ inputs=[document_input, summary_input, api_name_input, api_key_input, save_value],
281
+ outputs=output
282
+ )
283
+
284
+ return document_input, summary_input, api_name_input, api_key_input, evaluate_button, output
285
+
286
+
287
+ def parse_output(output: str, max: float) -> float:
288
+ """
289
+ Function that extracts numerical score from the beginning of string
290
+
291
+ Args:
292
+ output (str): String to search
293
+ max (float): Maximum score allowed
294
+
295
+ Returns:
296
+ float: The extracted score
297
+ """
298
+ matched: List[str] = re.findall(r"(?<!\S)\d+(?:\.\d+)?", output)
299
+ if matched:
300
+ if len(matched) == 1:
301
+ score = float(matched[0])
302
+ if score > max:
303
+ raise ValueError(f"Parsed number: {score} was larger than max score: {max}")
304
+ else:
305
+ raise ValueError(f"More than one number detected in input. Input to parser was: {output}")
306
+ else:
307
+ raise ValueError(f'No number detected in input. Input to parser was "{output}". ')
308
+ return score
309
+
310
+ def geval_summarization(
311
+ prompt_with_src_and_gen: str,
312
+ max_score: float,
313
+ api_endpoint: str,
314
+ api_key: str,
315
+ ) -> float:
316
+ model = get_model_from_config(api_endpoint)
317
+
318
+ try:
319
+ for attempt in Retrying(
320
+ reraise=True,
321
+ before_sleep=before_sleep_log(logger, logging.INFO),
322
+ after=after_log(logger, logging.INFO),
323
+ wait=wait_random_exponential(multiplier=1, min=1, max=120),
324
+ stop=stop_after_attempt(10),
325
+ ):
326
+ with attempt:
327
+ system_message="You are a helpful AI assistant"
328
+ # TEMP setting for Confabulation check
329
+ temp = 0.7
330
+ logging.info(f"Debug - geval_summarization Function - API Endpoint: {api_endpoint}")
331
+ try:
332
+ response = chat_api_call(api_endpoint, api_key, prompt_with_src_and_gen, "", temp, system_message)
333
+ except Exception as e:
334
+ raise ValueError(f"Unsupported API endpoint: {api_endpoint}")
335
+ except RetryError:
336
+ logger.exception(f"geval {api_endpoint} call failed\nInput prompt was: {prompt_with_src_and_gen}")
337
+ raise
338
+
339
+ try:
340
+ score = parse_output(response, max_score)
341
+ except ValueError as e:
342
+ logger.warning(f"Error parsing output: {e}")
343
+ score = 0
344
+
345
+ return score
346
+
347
+
348
+ def get_model_from_config(api_name: str) -> str:
349
+ model = config.get('models', api_name)
350
+ if isinstance(model, dict):
351
+ # If the model is a dictionary, return a specific key or a default value
352
+ return model.get('name', str(model)) # Adjust 'name' to the appropriate key if needed
353
+ return str(model) if model is not None else ""
354
+
355
+ def aggregate_llm_scores(llm_responses: List[str], max_score: float) -> float:
356
+ """Parse and average valid scores from the generated responses of
357
+ the G-Eval LLM call.
358
+
359
+ Args:
360
+ llm_responses (List[str]): List of scores from multiple LLMs
361
+ max_score (float): The maximum score allowed.
362
+
363
+ Returns:
364
+ float: The average of all the valid scores
365
+ """
366
+ all_scores = []
367
+ error_count = 0
368
+ for generated in llm_responses:
369
+ try:
370
+ parsed = parse_output(generated, max_score)
371
+ all_scores.append(parsed)
372
+ except ValueError as e:
373
+ logger.warning(e)
374
+ error_count += 1
375
+ if error_count:
376
+ logger.warning(f"{error_count} out of 20 scores were discarded due to corrupt g-eval generation")
377
+ score = sum(all_scores) / len(all_scores)
378
+ return score
379
+
380
+
381
+ def validate_inputs(document: str, summary: str, api_name: str, api_key: str) -> None:
382
+ """
383
+ Validate inputs for the G-Eval function.
384
+
385
+ Args:
386
+ document (str): The source document
387
+ summary (str): The summary to evaluate
388
+ api_name (str): The name of the API to use
389
+ api_key (str): The API key
390
+
391
+ Raises:
392
+ ValueError: If any of the inputs are invalid
393
+ """
394
+ if not document.strip():
395
+ raise ValueError("Source document cannot be empty")
396
+ if not summary.strip():
397
+ raise ValueError("Summary cannot be empty")
398
+ if api_name.lower() not in ["openai", "anthropic", "cohere", "groq", "openrouter", "deepseek", "huggingface",
399
+ "mistral", "llama.cpp", "kobold", "ooba", "tabbyapi", "vllm", "local-llm", "ollama"]:
400
+ raise ValueError(f"Unsupported API: {api_name}")
401
+
402
+
403
+ def detailed_api_error(api_name: str, error: Exception) -> str:
404
+ """
405
+ Generate a detailed error message for API failures.
406
+
407
+ Args:
408
+ api_name (str): The name of the API that failed
409
+ error (Exception): The exception that was raised
410
+
411
+ Returns:
412
+ str: A detailed error message
413
+ """
414
+ error_type = type(error).__name__
415
+ error_message = str(error)
416
+ return f"API Failure: {api_name}\nError Type: {error_type}\nError Message: {error_message}\nPlease check your API key and network connection, and try again."
417
+
418
+
419
+ def save_eval_results(results: Dict[str, Any], filename: str = "geval_results.json") -> None:
420
+ """
421
+ Save evaluation results to a JSON file.
422
+
423
+ Args:
424
+ results (Dict[str, Any]): The evaluation results
425
+ filename (str): The name of the file to save results to
426
+ """
427
+ with open(filename, 'w') as f:
428
+ json.dump(results, f, indent=2)
429
+ print(f"Results saved to {filename}")
430
+
431
+
432
+
433
+
434
+ #
435
+ #
436
+ #######################################################################################################################
437
+ #
438
+ # Taken from: https://github.com/microsoft/promptflow/blob/b5a68f45e4c3818a29e2f79a76f2e73b8ea6be44/src/promptflow-core/promptflow/_core/metric_logger.py
439
+
440
+ class MetricLoggerManager:
441
+ _instance = None
442
+
443
+ def __init__(self):
444
+ self._metric_loggers = []
445
+
446
+ @staticmethod
447
+ def get_instance() -> "MetricLoggerManager":
448
+ if MetricLoggerManager._instance is None:
449
+ MetricLoggerManager._instance = MetricLoggerManager()
450
+ return MetricLoggerManager._instance
451
+
452
+ def log_metric(self, key, value, variant_id=None):
453
+ for logger in self._metric_loggers:
454
+ if len(inspect.signature(logger).parameters) == 2:
455
+ logger(key, value) # If the logger only accepts two parameters, we don't pass variant_id
456
+ else:
457
+ logger(key, value, variant_id)
458
+
459
+ def add_metric_logger(self, logger_func: Callable):
460
+ existing_logger = next((logger for logger in self._metric_loggers if logger is logger_func), None)
461
+ if existing_logger:
462
+ return
463
+ if not callable(logger_func):
464
+ return
465
+ sign = inspect.signature(logger_func)
466
+ # We accept two kinds of metric loggers:
467
+ # def log_metric(k, v)
468
+ # def log_metric(k, v, variant_id)
469
+ if len(sign.parameters) not in [2, 3]:
470
+ return
471
+ self._metric_loggers.append(logger_func)
472
+
473
+ def remove_metric_logger(self, logger_func: Callable):
474
+ self._metric_loggers.remove(logger_func)
475
+
476
+
477
+ def log_metric(key, value, variant_id=None):
478
+ """Log a metric for current promptflow run.
479
+
480
+ :param key: Metric name.
481
+ :type key: str
482
+ :param value: Metric value.
483
+ :type value: float
484
+ :param variant_id: Variant id for the metric.
485
+ :type variant_id: str
486
+ """
487
+ MetricLoggerManager.get_instance().log_metric(key, value, variant_id)
488
+
489
+
490
+ def add_metric_logger(logger_func: Callable):
491
+ MetricLoggerManager.get_instance().add_metric_logger(logger_func)
492
+
493
+
494
+ def remove_metric_logger(logger_func: Callable):
495
+ MetricLoggerManager.get_instance().remove_metric_logger(logger_func)
496
+ #
497
+ # End of G-Eval.py
498
+ #######################################################################################################################
App_Function_Libraries/Books/.pytest_cache/.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ # Created by pytest automatically.
2
+ *
App_Function_Libraries/Books/.pytest_cache/CACHEDIR.TAG ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ Signature: 8a477f597d28d172789f06886806bc55
2
+ # This file is a cache directory tag created by pytest.
3
+ # For information about cache directory tags, see:
4
+ # https://bford.info/cachedir/spec.html
App_Function_Libraries/Books/.pytest_cache/README.md ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ # pytest cache directory #
2
+
3
+ This directory contains data from the pytest's cache plugin,
4
+ which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
5
+
6
+ **Do not** commit this to version control.
7
+
8
+ See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.
App_Function_Libraries/Books/.pytest_cache/v/cache/lastfailed ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_file": true,
3
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_missing_metadata": true,
4
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_with_auto_summarize": true,
5
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_process_zip_file": true,
6
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_epub_file": true,
7
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_zip_file": true,
8
+ "test_Book_Ingestion_lib.py": true,
9
+ "test_Book_Ingestion_tab.py": true
10
+ }
App_Function_Libraries/Books/.pytest_cache/v/cache/nodeids ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_file",
3
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_invalid_file",
4
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_missing_metadata",
5
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_with_auto_summarize",
6
+ "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_process_zip_file",
7
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_epub_file",
8
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_no_file",
9
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_unsupported_file",
10
+ "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_zip_file"
11
+ ]
App_Function_Libraries/Books/.pytest_cache/v/cache/stepwise ADDED
@@ -0,0 +1 @@
 
 
1
+ []
App_Function_Libraries/Books/Book_Ingestion_Lib.py ADDED
@@ -0,0 +1,577 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Book_Ingestion_Lib.py
2
+ #########################################
3
+ # Library to hold functions for ingesting book files.#
4
+ #
5
+ ####################
6
+ # Function List
7
+ #
8
+ # 1. ingest_text_file(file_path, title=None, author=None, keywords=None):
9
+ # 2.
10
+ #
11
+ #
12
+ ####################
13
+ #
14
+ # Imports
15
+ import os
16
+ import re
17
+ import tempfile
18
+ import zipfile
19
+ from datetime import datetime
20
+ import logging
21
+ #
22
+ # External Imports
23
+ import ebooklib
24
+ from bs4 import BeautifulSoup
25
+ from ebooklib import epub
26
+ #
27
+ # Import Local
28
+ from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords, add_media_to_database
29
+ from App_Function_Libraries.Summarization.Summarization_General_Lib import perform_summarization
30
+ from App_Function_Libraries.Chunk_Lib import chunk_ebook_by_chapters
31
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
32
+ #
33
+ #######################################################################################################################
34
+ # Function Definitions
35
+ #
36
+
37
+ def import_epub(file_path,
38
+ title=None,
39
+ author=None,
40
+ keywords=None,
41
+ custom_prompt=None,
42
+ system_prompt=None,
43
+ summary=None,
44
+ auto_summarize=False,
45
+ api_name=None,
46
+ api_key=None,
47
+ chunk_options=None,
48
+ custom_chapter_pattern=None
49
+ ):
50
+ """
51
+ Imports an EPUB file, extracts its content, chunks it, optionally summarizes it, and adds it to the database.
52
+
53
+ Parameters:
54
+ - file_path (str): Path to the EPUB file.
55
+ - title (str, optional): Title of the book.
56
+ - author (str, optional): Author of the book.
57
+ - keywords (str, optional): Comma-separated keywords for the book.
58
+ - custom_prompt (str, optional): Custom user prompt for summarization.
59
+ - summary (str, optional): Predefined summary of the book.
60
+ - auto_summarize (bool, optional): Whether to auto-summarize the chunks.
61
+ - api_name (str, optional): API name for summarization.
62
+ - api_key (str, optional): API key for summarization.
63
+ - chunk_options (dict, optional): Options for chunking.
64
+ - custom_chapter_pattern (str, optional): Custom regex pattern for chapter detection.
65
+
66
+ Returns:
67
+ - str: Status message indicating success or failure.
68
+ """
69
+ try:
70
+ logging.info(f"Importing EPUB file from {file_path}")
71
+ log_counter("epub_import_attempt", labels={"file_path": file_path})
72
+
73
+ start_time = datetime.now()
74
+
75
+ # Convert EPUB to Markdown
76
+ markdown_content = epub_to_markdown(file_path)
77
+ logging.debug("Converted EPUB to Markdown.")
78
+
79
+ # Extract metadata if not provided
80
+ if not title or not author:
81
+ extracted_title, extracted_author = extract_epub_metadata(markdown_content)
82
+ title = title or extracted_title or os.path.splitext(os.path.basename(file_path))[0]
83
+ author = author or extracted_author or "Unknown"
84
+ logging.debug(f"Extracted metadata - Title: {title}, Author: {author}")
85
+
86
+ # Process keywords
87
+ keyword_list = [kw.strip() for kw in keywords.split(',')] if keywords else []
88
+ logging.debug(f"Keywords: {keyword_list}")
89
+
90
+ # Set default chunk options if not provided
91
+ if chunk_options is None:
92
+ chunk_options = {
93
+ 'method': 'chapter',
94
+ 'max_size': 500,
95
+ 'overlap': 200,
96
+ 'custom_chapter_pattern': custom_chapter_pattern
97
+ }
98
+ else:
99
+ # Ensure 'method' is set to 'chapter' when using chapter chunking
100
+ chunk_options.setdefault('method', 'chapter')
101
+ chunk_options.setdefault('custom_chapter_pattern', custom_chapter_pattern)
102
+
103
+ # Chunk the content by chapters
104
+ chunks = chunk_ebook_by_chapters(markdown_content, chunk_options)
105
+ logging.info(f"Total chunks created: {len(chunks)}")
106
+ log_histogram("epub_chunks_created", len(chunks), labels={"file_path": file_path})
107
+
108
+ if chunks:
109
+ logging.debug(f"Structure of first chunk: {chunks[0].keys()}")
110
+
111
+ # Handle summarization if enabled
112
+ if auto_summarize and api_name and api_key:
113
+ logging.info("Auto-summarization is enabled.")
114
+ summarized_chunks = []
115
+ for chunk in chunks:
116
+ chunk_text = chunk.get('text', '')
117
+ if chunk_text:
118
+ summary_text = perform_summarization(api_name, chunk_text, custom_prompt, api_key,
119
+ recursive_summarization=False, temp=None,
120
+ system_message=system_prompt
121
+ )
122
+ chunk['metadata']['summary'] = summary_text
123
+ summarized_chunks.append(chunk)
124
+ chunks = summarized_chunks
125
+ logging.info("Summarization of chunks completed.")
126
+ log_counter("epub_chunks_summarized", value=len(chunks), labels={"file_path": file_path})
127
+ else:
128
+ # If not summarizing, set a default summary or use provided summary
129
+ if summary:
130
+ logging.debug("Using provided summary.")
131
+ else:
132
+ summary = "No summary provided."
133
+
134
+ # Create info_dict
135
+ info_dict = {
136
+ 'title': title,
137
+ 'uploader': author,
138
+ 'ingestion_date': datetime.now().strftime('%Y-%m-%d')
139
+ }
140
+
141
+ # Prepare segments for database
142
+ segments = [{'Text': chunk.get('text', chunk.get('content', ''))} for chunk in chunks]
143
+ logging.debug(f"Prepared segments for database. Number of segments: {len(segments)}")
144
+
145
+ # Add to database
146
+ result = add_media_to_database(
147
+ url=file_path,
148
+ info_dict=info_dict,
149
+ segments=segments,
150
+ summary=summary,
151
+ keywords=keyword_list,
152
+ custom_prompt_input=custom_prompt,
153
+ whisper_model="Imported",
154
+ media_type="ebook",
155
+ overwrite=False
156
+ )
157
+
158
+ end_time = datetime.now()
159
+ processing_time = (end_time - start_time).total_seconds()
160
+ log_histogram("epub_import_duration", processing_time, labels={"file_path": file_path})
161
+
162
+ logging.info(f"Ebook '{title}' by {author} imported successfully. Database result: {result}")
163
+ log_counter("epub ingested into the DB successfully", labels={"file_path": file_path})
164
+ return f"Ebook '{title}' by {author} imported successfully. Database result: {result}"
165
+
166
+ except Exception as e:
167
+ logging.exception(f"Error importing ebook: {str(e)}")
168
+ log_counter("epub_import_error", labels={"file_path": file_path, "error": str(e)})
169
+ return f"Error importing ebook: {str(e)}"
170
+
171
+
172
+ # FIXME
173
+ def process_zip_file(zip_file,
174
+ title,
175
+ author,
176
+ keywords,
177
+ custom_prompt,
178
+ system_prompt,
179
+ summary,
180
+ auto_summarize,
181
+ api_name,
182
+ api_key,
183
+ chunk_options
184
+ ):
185
+ """
186
+ Processes a ZIP file containing multiple EPUB files and imports each one.
187
+
188
+ Parameters:
189
+ - zip_file (file-like object): The ZIP file to process.
190
+ - title (str): Title prefix for the books.
191
+ - author (str): Author name for the books.
192
+ - keywords (str): Comma-separated keywords.
193
+ - custom_prompt (str): Custom user prompt for summarization.
194
+ - summary (str): Predefined summary (not used in this context).
195
+ - auto_summarize (bool): Whether to auto-summarize the chunks.
196
+ - api_name (str): API name for summarization.
197
+ - api_key (str): API key for summarization.
198
+ - chunk_options (dict): Options for chunking.
199
+
200
+ Returns:
201
+ - str: Combined status messages for all EPUB files in the ZIP.
202
+ """
203
+ results = []
204
+ try:
205
+ with tempfile.TemporaryDirectory() as temp_dir:
206
+ zip_path = zip_file.name if hasattr(zip_file, 'name') else zip_file.path
207
+ logging.info(f"Extracting ZIP file {zip_path} to temporary directory {temp_dir}")
208
+ log_counter("zip_processing_attempt", labels={"zip_path": zip_path})
209
+
210
+ with zipfile.ZipFile(zip_path, 'r') as zip_ref:
211
+ zip_ref.extractall(temp_dir)
212
+
213
+ epub_files = [f for f in os.listdir(temp_dir) if f.lower().endswith('.epub')]
214
+ log_histogram("epub_files_in_zip", len(epub_files), labels={"zip_path": zip_path})
215
+
216
+ for filename in epub_files:
217
+ file_path = os.path.join(temp_dir, filename)
218
+ logging.info(f"Processing EPUB file {filename} from ZIP.")
219
+ result = import_epub(
220
+ file_path=file_path,
221
+ title=title,
222
+ author=author,
223
+ keywords=keywords,
224
+ custom_prompt=custom_prompt,
225
+ summary=summary,
226
+ auto_summarize=auto_summarize,
227
+ api_name=api_name,
228
+ api_key=api_key,
229
+ chunk_options=chunk_options,
230
+ custom_chapter_pattern=chunk_options.get('custom_chapter_pattern') if chunk_options else None
231
+ )
232
+ results.append(f"File: {filename} - {result}")
233
+
234
+ logging.info("Completed processing all EPUB files in the ZIP.")
235
+ log_counter("zip_processing_success", labels={"zip_path": zip_path})
236
+ except Exception as e:
237
+ logging.exception(f"Error processing ZIP file: {str(e)}")
238
+ log_counter("zip_processing_error", labels={"zip_path": zip_path, "error": str(e)})
239
+ return f"Error processing ZIP file: {str(e)}"
240
+
241
+ return "\n".join(results)
242
+
243
+
244
+ def import_file_handler(file,
245
+ title,
246
+ author,
247
+ keywords,
248
+ system_prompt,
249
+ custom_prompt,
250
+ auto_summarize,
251
+ api_name,
252
+ api_key,
253
+ max_chunk_size,
254
+ chunk_overlap,
255
+ custom_chapter_pattern
256
+ ):
257
+ try:
258
+ log_counter("file_import_attempt", labels={"file_name": file.name})
259
+
260
+ # Handle max_chunk_size
261
+ if isinstance(max_chunk_size, str):
262
+ max_chunk_size = int(max_chunk_size) if max_chunk_size.strip() else 4000
263
+ elif not isinstance(max_chunk_size, int):
264
+ max_chunk_size = 4000 # Default value if not a string or int
265
+
266
+ # Handle chunk_overlap
267
+ if isinstance(chunk_overlap, str):
268
+ chunk_overlap = int(chunk_overlap) if chunk_overlap.strip() else 0
269
+ elif not isinstance(chunk_overlap, int):
270
+ chunk_overlap = 0 # Default value if not a string or int
271
+
272
+ chunk_options = {
273
+ 'method': 'chapter',
274
+ 'max_size': max_chunk_size,
275
+ 'overlap': chunk_overlap,
276
+ 'custom_chapter_pattern': custom_chapter_pattern if custom_chapter_pattern else None
277
+ }
278
+
279
+ if file is None:
280
+ log_counter("file_import_error", labels={"error": "No file uploaded"})
281
+ return "No file uploaded."
282
+
283
+ file_path = file.name
284
+ if not os.path.exists(file_path):
285
+ log_counter("file_import_error", labels={"error": "File not found", "file_name": file.name})
286
+ return "Uploaded file not found."
287
+
288
+ start_time = datetime.now()
289
+
290
+ if file_path.lower().endswith('.epub'):
291
+ status = import_epub(
292
+ file_path,
293
+ title,
294
+ author,
295
+ keywords,
296
+ custom_prompt=custom_prompt,
297
+ system_prompt=system_prompt,
298
+ summary=None,
299
+ auto_summarize=auto_summarize,
300
+ api_name=api_name,
301
+ api_key=api_key,
302
+ chunk_options=chunk_options,
303
+ custom_chapter_pattern=custom_chapter_pattern
304
+ )
305
+ log_counter("epub_import_success", labels={"file_name": file.name})
306
+ result = f"📚 EPUB Imported Successfully:\n{status}"
307
+ elif file.name.lower().endswith('.zip'):
308
+ status = process_zip_file(
309
+ zip_file=file,
310
+ title=title,
311
+ author=author,
312
+ keywords=keywords,
313
+ custom_prompt=custom_prompt,
314
+ system_prompt=system_prompt,
315
+ summary=None,
316
+ auto_summarize=auto_summarize,
317
+ api_name=api_name,
318
+ api_key=api_key,
319
+ chunk_options=chunk_options
320
+ )
321
+ log_counter("zip_import_success", labels={"file_name": file.name})
322
+ result = f"📦 ZIP Processed Successfully:\n{status}"
323
+ elif file.name.lower().endswith(('.chm', '.html', '.pdf', '.xml', '.opml')):
324
+ file_type = file.name.split('.')[-1].upper()
325
+ log_counter("unsupported_file_type", labels={"file_type": file_type})
326
+ result = f"{file_type} file import is not yet supported."
327
+ else:
328
+ log_counter("unsupported_file_type", labels={"file_type": file.name.split('.')[-1]})
329
+ result = "❌ Unsupported file type. Please upload an `.epub` file or a `.zip` file containing `.epub` files."
330
+
331
+ end_time = datetime.now()
332
+ processing_time = (end_time - start_time).total_seconds()
333
+ log_histogram("file_import_duration", processing_time, labels={"file_name": file.name})
334
+
335
+ return result
336
+
337
+ except ValueError as ve:
338
+ logging.exception(f"Error parsing input values: {str(ve)}")
339
+ log_counter("file_import_error", labels={"error": "Invalid input", "file_name": file.name})
340
+ return f"❌ Error: Invalid input for chunk size or overlap. Please enter valid numbers."
341
+ except Exception as e:
342
+ logging.exception(f"Error during file import: {str(e)}")
343
+ log_counter("file_import_error", labels={"error": str(e), "file_name": file.name})
344
+ return f"❌ Error during import: {str(e)}"
345
+
346
+
347
+ def read_epub(file_path):
348
+ """
349
+ Reads and extracts text from an EPUB file.
350
+
351
+ Parameters:
352
+ - file_path (str): Path to the EPUB file.
353
+
354
+ Returns:
355
+ - str: Extracted text content from the EPUB.
356
+ """
357
+ try:
358
+ logging.info(f"Reading EPUB file from {file_path}")
359
+ book = epub.read_epub(file_path)
360
+ chapters = []
361
+ for item in book.get_items():
362
+ if item.get_type() == ebooklib.ITEM_DOCUMENT:
363
+ chapters.append(item.get_content())
364
+
365
+ text = ""
366
+ for html_content in chapters:
367
+ soup = BeautifulSoup(html_content, 'html.parser')
368
+ text += soup.get_text(separator='\n\n') + "\n\n"
369
+ logging.debug("EPUB content extraction completed.")
370
+ return text
371
+ except Exception as e:
372
+ logging.exception(f"Error reading EPUB file: {str(e)}")
373
+ raise
374
+
375
+
376
+ # Ingest a text file into the database with Title/Author/Keywords
377
+ def extract_epub_metadata(content):
378
+ title_match = re.search(r'Title:\s*(.*?)\n', content)
379
+ author_match = re.search(r'Author:\s*(.*?)\n', content)
380
+
381
+ title = title_match.group(1) if title_match else None
382
+ author = author_match.group(1) if author_match else None
383
+
384
+ return title, author
385
+
386
+
387
+ def ingest_text_file(file_path, title=None, author=None, keywords=None):
388
+ """
389
+ Ingests a plain text file into the database with optional metadata.
390
+
391
+ Parameters:
392
+ - file_path (str): Path to the text file.
393
+ - title (str, optional): Title of the document.
394
+ - author (str, optional): Author of the document.
395
+ - keywords (str, optional): Comma-separated keywords.
396
+
397
+ Returns:
398
+ - str: Status message indicating success or failure.
399
+ """
400
+ try:
401
+ with open(file_path, 'r', encoding='utf-8') as file:
402
+ content = file.read()
403
+
404
+ # Check if it's a converted epub and extract metadata if so
405
+ if 'epub_converted' in (keywords or '').lower():
406
+ extracted_title, extracted_author = extract_epub_metadata(content)
407
+ title = title or extracted_title
408
+ author = author or extracted_author
409
+ logging.debug(f"Extracted metadata for converted EPUB - Title: {title}, Author: {author}")
410
+
411
+ # If title is still not provided, use the filename without extension
412
+ if not title:
413
+ title = os.path.splitext(os.path.basename(file_path))[0]
414
+
415
+ # If author is still not provided, set it to 'Unknown'
416
+ if not author:
417
+ author = 'Unknown'
418
+
419
+ # If keywords are not provided, use a default keyword
420
+ if not keywords:
421
+ keywords = 'text_file,epub_converted'
422
+ else:
423
+ keywords = f'text_file,epub_converted,{keywords}'
424
+
425
+ # Add the text file to the database
426
+ add_media_with_keywords(
427
+ url=file_path,
428
+ title=title,
429
+ media_type='document',
430
+ content=content,
431
+ keywords=keywords,
432
+ prompt='No prompt for text files',
433
+ summary='No summary for text files',
434
+ transcription_model='None',
435
+ author=author,
436
+ ingestion_date=datetime.now().strftime('%Y-%m-%d')
437
+ )
438
+
439
+ logging.info(f"Text file '{title}' by {author} ingested successfully.")
440
+ return f"Text file '{title}' by {author} ingested successfully."
441
+ except Exception as e:
442
+ logging.error(f"Error ingesting text file: {str(e)}")
443
+ return f"Error ingesting text file: {str(e)}"
444
+
445
+
446
+ def ingest_folder(folder_path, keywords=None):
447
+ """
448
+ Ingests all text files within a specified folder.
449
+
450
+ Parameters:
451
+ - folder_path (str): Path to the folder containing text files.
452
+ - keywords (str, optional): Comma-separated keywords to add to each file.
453
+
454
+ Returns:
455
+ - str: Combined status messages for all ingested text files.
456
+ """
457
+ results = []
458
+ try:
459
+ logging.info(f"Ingesting all text files from folder {folder_path}")
460
+ for filename in os.listdir(folder_path):
461
+ if filename.lower().endswith('.txt'):
462
+ file_path = os.path.join(folder_path, filename)
463
+ result = ingest_text_file(file_path, keywords=keywords)
464
+ results.append(result)
465
+ logging.info("Completed ingestion of all text files in the folder.")
466
+ except Exception as e:
467
+ logging.exception(f"Error ingesting folder: {str(e)}")
468
+ return f"Error ingesting folder: {str(e)}"
469
+
470
+ return "\n".join(results)
471
+
472
+
473
+ def epub_to_markdown(epub_path):
474
+ """
475
+ Converts an EPUB file to Markdown format, including the table of contents and chapter contents.
476
+
477
+ Parameters:
478
+ - epub_path (str): Path to the EPUB file.
479
+
480
+ Returns:
481
+ - str: Markdown-formatted content of the EPUB.
482
+ """
483
+ try:
484
+ logging.info(f"Converting EPUB to Markdown from {epub_path}")
485
+ book = epub.read_epub(epub_path)
486
+ markdown_content = "# Table of Contents\n\n"
487
+ chapters = []
488
+
489
+ # Extract and format the table of contents
490
+ toc = book.toc
491
+ for item in toc:
492
+ if isinstance(item, tuple):
493
+ section, children = item
494
+ level = 1
495
+ markdown_content += format_toc_item(section, level)
496
+ for child in children:
497
+ markdown_content += format_toc_item(child, level + 1)
498
+ else:
499
+ markdown_content += format_toc_item(item, 1)
500
+
501
+ markdown_content += "\n---\n\n"
502
+
503
+ # Process each chapter
504
+ for item in book.get_items():
505
+ if item.get_type() == ebooklib.ITEM_DOCUMENT:
506
+ chapter_content = item.get_content().decode('utf-8')
507
+ soup = BeautifulSoup(chapter_content, 'html.parser')
508
+
509
+ # Extract chapter title
510
+ title = soup.find(['h1', 'h2', 'h3'])
511
+ if title:
512
+ chapter_title = title.get_text()
513
+ markdown_content += f"# {chapter_title}\n\n"
514
+
515
+ # Process chapter content
516
+ for elem in soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'ul', 'ol']):
517
+ if elem.name.startswith('h'):
518
+ level = int(elem.name[1])
519
+ markdown_content += f"{'#' * level} {elem.get_text()}\n\n"
520
+ elif elem.name == 'p':
521
+ markdown_content += f"{elem.get_text()}\n\n"
522
+ elif elem.name in ['ul', 'ol']:
523
+ for li in elem.find_all('li'):
524
+ prefix = '-' if elem.name == 'ul' else '1.'
525
+ markdown_content += f"{prefix} {li.get_text()}\n"
526
+ markdown_content += "\n"
527
+
528
+ markdown_content += "---\n\n"
529
+
530
+ logging.debug("EPUB to Markdown conversion completed.")
531
+ return markdown_content
532
+
533
+ except Exception as e:
534
+ logging.exception(f"Error converting EPUB to Markdown: {str(e)}")
535
+ raise
536
+
537
+
538
+ def format_toc_item(item, level):
539
+ """
540
+ Formats a table of contents item into Markdown list format.
541
+
542
+ Parameters:
543
+ - item (epub.Link or epub.Section): TOC item.
544
+ - level (int): Heading level for indentation.
545
+
546
+ Returns:
547
+ - str: Markdown-formatted TOC item.
548
+ """
549
+ try:
550
+ if isinstance(item, epub.Link):
551
+ title = item.title
552
+ elif isinstance(item, epub.Section):
553
+ title = item.title
554
+ else:
555
+ title = str(item)
556
+
557
+ return f"{' ' * (level - 1)}- [{title}](#{slugify(title)})\n"
558
+ except Exception as e:
559
+ logging.exception(f"Error formatting TOC item: {str(e)}")
560
+ return ""
561
+
562
+
563
+ def slugify(text):
564
+ """
565
+ Converts a string into a slug suitable for Markdown links.
566
+
567
+ Parameters:
568
+ - text (str): The text to slugify.
569
+
570
+ Returns:
571
+ - str: Slugified text.
572
+ """
573
+ return re.sub(r'[\W_]+', '-', text.lower()).strip('-')
574
+
575
+ #
576
+ # End of Function Definitions
577
+ #######################################################################################################################
App_Function_Libraries/Books/__init__.py ADDED
File without changes
App_Function_Libraries/Character_Chat/Character_Chat_Lib.py ADDED
@@ -0,0 +1,607 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Character_Chat_Lib.py
2
+ # Description: Functions for character chat cards.
3
+ #
4
+ # Imports
5
+ import json
6
+ import logging
7
+ import io
8
+ import base64
9
+ import time
10
+ from typing import Dict, Any, Optional, List, Tuple
11
+ #
12
+ # External Imports
13
+ from PIL import Image
14
+ #
15
+ # Local imports
16
+ from App_Function_Libraries.DB.DB_Manager import get_character_card_by_id, get_character_chat_by_id
17
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
18
+ #
19
+ # Constants
20
+ ####################################################################################################
21
+ #
22
+ # Functions
23
+
24
+ # Using https://github.com/malfoyslastname/character-card-spec-v2 as the standard for v2 character cards
25
+
26
+ #################################################################################
27
+ #
28
+ # Placeholder functions:
29
+
30
+ def replace_placeholders(text: str, char_name: str, user_name: str) -> str:
31
+ """
32
+ Replace placeholders in the given text with appropriate values.
33
+
34
+ Args:
35
+ text (str): The text containing placeholders.
36
+ char_name (str): The name of the character.
37
+ user_name (str): The name of the user.
38
+
39
+ Returns:
40
+ str: The text with placeholders replaced.
41
+ """
42
+ replacements = {
43
+ '{{char}}': char_name,
44
+ '{{user}}': user_name,
45
+ '{{random_user}}': user_name # Assuming random_user is the same as user for simplicity
46
+ }
47
+
48
+ for placeholder, value in replacements.items():
49
+ text = text.replace(placeholder, value)
50
+
51
+ return text
52
+
53
+ def replace_user_placeholder(history, user_name):
54
+ """
55
+ Replaces all instances of '{{user}}' in the chat history with the actual user name.
56
+
57
+ Args:
58
+ history (list): The current chat history as a list of tuples (user_message, bot_message).
59
+ user_name (str): The name entered by the user.
60
+
61
+ Returns:
62
+ list: Updated chat history with placeholders replaced.
63
+ """
64
+ if not user_name:
65
+ user_name = "User" # Default name if none provided
66
+
67
+ updated_history = []
68
+ for user_msg, bot_msg in history:
69
+ # Replace in user message
70
+ if user_msg:
71
+ user_msg = user_msg.replace("{{user}}", user_name)
72
+ # Replace in bot message
73
+ if bot_msg:
74
+ bot_msg = bot_msg.replace("{{user}}", user_name)
75
+ updated_history.append((user_msg, bot_msg))
76
+ return updated_history
77
+
78
+ #
79
+ # End of Placeholder functions
80
+ #################################################################################
81
+
82
+ #################################################################################
83
+ #
84
+ # Functions for character card processing:
85
+
86
+ def extract_character_id(choice: str) -> int:
87
+ """Extract the character ID from the dropdown selection string."""
88
+ log_counter("extract_character_id_attempt")
89
+ try:
90
+ character_id = int(choice.split('(ID: ')[1].rstrip(')'))
91
+ log_counter("extract_character_id_success")
92
+ return character_id
93
+ except Exception as e:
94
+ log_counter("extract_character_id_error", labels={"error": str(e)})
95
+ raise
96
+
97
+ def load_character_wrapper(character_id: int, user_name: str) -> Tuple[Dict[str, Any], List[Tuple[Optional[str], str]], Optional[Image.Image]]:
98
+ """Wrapper function to load character and image using the extracted ID."""
99
+ log_counter("load_character_wrapper_attempt")
100
+ start_time = time.time()
101
+ try:
102
+ char_data, chat_history, img = load_character_and_image(character_id, user_name)
103
+ load_duration = time.time() - start_time
104
+ log_histogram("load_character_wrapper_duration", load_duration)
105
+ log_counter("load_character_wrapper_success")
106
+ return char_data, chat_history, img
107
+ except Exception as e:
108
+ log_counter("load_character_wrapper_error", labels={"error": str(e)})
109
+ raise
110
+
111
+ def parse_character_book(book_data: Dict[str, Any]) -> Dict[str, Any]:
112
+ """
113
+ Parse the character book data from a V2 character card.
114
+
115
+ Args:
116
+ book_data (Dict[str, Any]): The raw character book data from the character card.
117
+
118
+ Returns:
119
+ Dict[str, Any]: The parsed and structured character book data.
120
+ """
121
+ parsed_book = {
122
+ 'name': book_data.get('name', ''),
123
+ 'description': book_data.get('description', ''),
124
+ 'scan_depth': book_data.get('scan_depth'),
125
+ 'token_budget': book_data.get('token_budget'),
126
+ 'recursive_scanning': book_data.get('recursive_scanning', False),
127
+ 'extensions': book_data.get('extensions', {}),
128
+ 'entries': []
129
+ }
130
+
131
+ for entry in book_data.get('entries', []):
132
+ parsed_entry = {
133
+ 'keys': entry['keys'],
134
+ 'content': entry['content'],
135
+ 'extensions': entry.get('extensions', {}),
136
+ 'enabled': entry['enabled'],
137
+ 'insertion_order': entry['insertion_order'],
138
+ 'case_sensitive': entry.get('case_sensitive', False),
139
+ 'name': entry.get('name', ''),
140
+ 'priority': entry.get('priority'),
141
+ 'id': entry.get('id'),
142
+ 'comment': entry.get('comment', ''),
143
+ 'selective': entry.get('selective', False),
144
+ 'secondary_keys': entry.get('secondary_keys', []),
145
+ 'constant': entry.get('constant', False),
146
+ 'position': entry.get('position')
147
+ }
148
+ parsed_book['entries'].append(parsed_entry)
149
+
150
+ return parsed_book
151
+
152
+ def load_character_and_image(character_id: int, user_name: str) -> Tuple[Optional[Dict[str, Any]], List[Tuple[Optional[str], str]], Optional[Image.Image]]:
153
+ """
154
+ Load a character and its associated image based on the character ID.
155
+
156
+ Args:
157
+ character_id (int): The ID of the character to load.
158
+ user_name (str): The name of the user, used for placeholder replacement.
159
+
160
+ Returns:
161
+ Tuple[Optional[Dict[str, Any]], List[Tuple[Optional[str], str]], Optional[Image.Image]]:
162
+ A tuple containing the character data, chat history, and character image (if available).
163
+ """
164
+ log_counter("load_character_and_image_attempt")
165
+ start_time = time.time()
166
+ try:
167
+ char_data = get_character_card_by_id(character_id)
168
+ if not char_data:
169
+ log_counter("load_character_and_image_no_data")
170
+ logging.warning(f"No character data found for ID: {character_id}")
171
+ return None, [], None
172
+
173
+ # Replace placeholders in character data
174
+ for field in ['first_mes', 'mes_example', 'scenario', 'description', 'personality']:
175
+ if field in char_data:
176
+ char_data[field] = replace_placeholders(char_data[field], char_data['name'], user_name)
177
+
178
+ # Replace placeholders in first_mes
179
+ first_mes = char_data.get('first_mes', "Hello! I'm ready to chat.")
180
+ first_mes = replace_placeholders(first_mes, char_data['name'], user_name)
181
+
182
+ chat_history = [(None, first_mes)] if first_mes else []
183
+
184
+ img = None
185
+ if char_data.get('image'):
186
+ try:
187
+ image_data = base64.b64decode(char_data['image'])
188
+ img = Image.open(io.BytesIO(image_data)).convert("RGBA")
189
+ log_counter("load_character_image_success")
190
+ except Exception as e:
191
+ log_counter("load_character_image_error", labels={"error": str(e)})
192
+ logging.error(f"Error processing image for character '{char_data['name']}': {e}")
193
+
194
+ load_duration = time.time() - start_time
195
+ log_histogram("load_character_and_image_duration", load_duration)
196
+ log_counter("load_character_and_image_success")
197
+ return char_data, chat_history, img
198
+
199
+ except Exception as e:
200
+ log_counter("load_character_and_image_error", labels={"error": str(e)})
201
+ logging.error(f"Error in load_character_and_image: {e}")
202
+ return None, [], None
203
+
204
+ def load_chat_and_character(chat_id: int, user_name: str) -> Tuple[Optional[Dict[str, Any]], List[Tuple[str, str]], Optional[Image.Image]]:
205
+ """
206
+ Load a chat and its associated character, including the character image and process templates.
207
+
208
+ Args:
209
+ chat_id (int): The ID of the chat to load.
210
+ user_name (str): The name of the user.
211
+
212
+ Returns:
213
+ Tuple[Optional[Dict[str, Any]], List[Tuple[str, str]], Optional[Image.Image]]:
214
+ A tuple containing the character data, processed chat history, and character image (if available).
215
+ """
216
+ log_counter("load_chat_and_character_attempt")
217
+ start_time = time.time()
218
+ try:
219
+ # Load the chat
220
+ chat = get_character_chat_by_id(chat_id)
221
+ if not chat:
222
+ log_counter("load_chat_and_character_no_chat")
223
+ logging.warning(f"No chat found with ID: {chat_id}")
224
+ return None, [], None
225
+
226
+ # Load the associated character
227
+ character_id = chat['character_id']
228
+ char_data = get_character_card_by_id(character_id)
229
+ if not char_data:
230
+ log_counter("load_chat_and_character_no_character")
231
+ logging.warning(f"No character found for chat ID: {chat_id}")
232
+ return None, chat['chat_history'], None
233
+
234
+ # Process the chat history
235
+ processed_history = process_chat_history(chat['chat_history'], char_data['name'], user_name)
236
+
237
+ # Load the character image
238
+ img = None
239
+ if char_data.get('image'):
240
+ try:
241
+ image_data = base64.b64decode(char_data['image'])
242
+ img = Image.open(io.BytesIO(image_data)).convert("RGBA")
243
+ log_counter("load_chat_character_image_success")
244
+ except Exception as e:
245
+ log_counter("load_chat_character_image_error", labels={"error": str(e)})
246
+ logging.error(f"Error processing image for character '{char_data['name']}': {e}")
247
+
248
+ # Process character data templates
249
+ for field in ['first_mes', 'mes_example', 'scenario', 'description', 'personality']:
250
+ if field in char_data:
251
+ char_data[field] = replace_placeholders(char_data[field], char_data['name'], user_name)
252
+
253
+ load_duration = time.time() - start_time
254
+ log_histogram("load_chat_and_character_duration", load_duration)
255
+ log_counter("load_chat_and_character_success")
256
+ return char_data, processed_history, img
257
+
258
+ except Exception as e:
259
+ log_counter("load_chat_and_character_error", labels={"error": str(e)})
260
+ logging.error(f"Error in load_chat_and_character: {e}")
261
+ return None, [], None
262
+
263
+
264
+ def extract_json_from_image(image_file):
265
+ logging.debug(f"Attempting to extract JSON from image: {image_file.name}")
266
+ log_counter("extract_json_from_image_attempt")
267
+ start_time = time.time()
268
+ try:
269
+ with Image.open(image_file) as img:
270
+ logging.debug("Image opened successfully")
271
+ metadata = img.info
272
+ if 'chara' in metadata:
273
+ logging.debug("Found 'chara' in image metadata")
274
+ chara_content = metadata['chara']
275
+ logging.debug(f"Content of 'chara' metadata (first 100 chars): {chara_content[:100]}...")
276
+ try:
277
+ decoded_content = base64.b64decode(chara_content).decode('utf-8')
278
+ logging.debug(f"Decoded content (first 100 chars): {decoded_content[:100]}...")
279
+ log_counter("extract_json_from_image_metadata_success")
280
+ return decoded_content
281
+ except Exception as e:
282
+ logging.error(f"Error decoding base64 content: {e}")
283
+ log_counter("extract_json_from_image_decode_error", labels={"error": str(e)})
284
+
285
+ logging.warning("'chara' not found in metadata, attempting to find JSON data in image bytes")
286
+ # Alternative method to extract embedded JSON from image bytes if metadata is not available
287
+ img_byte_arr = io.BytesIO()
288
+ img.save(img_byte_arr, format='PNG')
289
+ img_bytes = img_byte_arr.getvalue()
290
+ img_str = img_bytes.decode('latin1')
291
+
292
+ # Search for JSON-like structures in the image bytes
293
+ json_start = img_str.find('{')
294
+ json_end = img_str.rfind('}')
295
+ if json_start != -1 and json_end != -1 and json_end > json_start:
296
+ possible_json = img_str[json_start:json_end+1]
297
+ try:
298
+ json.loads(possible_json)
299
+ logging.debug("Found JSON data in image bytes")
300
+ log_counter("extract_json_from_image_bytes_success")
301
+ return possible_json
302
+ except json.JSONDecodeError:
303
+ logging.debug("No valid JSON found in image bytes")
304
+ log_counter("extract_json_from_image_invalid_json")
305
+
306
+ logging.warning("No JSON data found in the image")
307
+ log_counter("extract_json_from_image_no_json_found")
308
+ except Exception as e:
309
+ log_counter("extract_json_from_image_error", labels={"error": str(e)})
310
+ logging.error(f"Error extracting JSON from image: {e}")
311
+
312
+ extract_duration = time.time() - start_time
313
+ log_histogram("extract_json_from_image_duration", extract_duration)
314
+ return None
315
+
316
+
317
+ def load_chat_history(file):
318
+ log_counter("load_chat_history_attempt")
319
+ start_time = time.time()
320
+ try:
321
+ content = file.read().decode('utf-8')
322
+ chat_data = json.loads(content)
323
+
324
+ # Extract history and character name from the loaded data
325
+ history = chat_data.get('history') or chat_data.get('messages')
326
+ character_name = chat_data.get('character') or chat_data.get('character_name')
327
+
328
+ if not history or not character_name:
329
+ log_counter("load_chat_history_incomplete_data")
330
+ logging.error("Chat history or character name missing in the imported file.")
331
+ return None, None
332
+
333
+ load_duration = time.time() - start_time
334
+ log_histogram("load_chat_history_duration", load_duration)
335
+ log_counter("load_chat_history_success")
336
+ return history, character_name
337
+ except Exception as e:
338
+ log_counter("load_chat_history_error", labels={"error": str(e)})
339
+ logging.error(f"Error loading chat history: {e}")
340
+ return None, None
341
+
342
+
343
+ def process_chat_history(chat_history: List[Tuple[str, str]], char_name: str, user_name: str) -> List[Tuple[str, str]]:
344
+ """
345
+ Process the chat history to replace placeholders in both user and character messages.
346
+
347
+ Args:
348
+ chat_history (List[Tuple[str, str]]): The chat history.
349
+ char_name (str): The name of the character.
350
+ user_name (str): The name of the user.
351
+
352
+ Returns:
353
+ List[Tuple[str, str]]: The processed chat history.
354
+ """
355
+ log_counter("process_chat_history_attempt")
356
+ start_time = time.time()
357
+ try:
358
+ processed_history = []
359
+ for user_msg, char_msg in chat_history:
360
+ if user_msg:
361
+ user_msg = replace_placeholders(user_msg, char_name, user_name)
362
+ if char_msg:
363
+ char_msg = replace_placeholders(char_msg, char_name, user_name)
364
+ processed_history.append((user_msg, char_msg))
365
+
366
+ process_duration = time.time() - start_time
367
+ log_histogram("process_chat_history_duration", process_duration)
368
+ log_counter("process_chat_history_success", labels={"message_count": len(chat_history)})
369
+ return processed_history
370
+ except Exception as e:
371
+ log_counter("process_chat_history_error", labels={"error": str(e)})
372
+ logging.error(f"Error processing chat history: {e}")
373
+ raise
374
+
375
+ def validate_character_book(book_data):
376
+ """
377
+ Validate the 'character_book' field in the character card.
378
+
379
+ Args:
380
+ book_data (dict): The character book data.
381
+
382
+ Returns:
383
+ Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages.
384
+ """
385
+ validation_messages = []
386
+
387
+ # Optional fields with expected types
388
+ optional_fields = {
389
+ 'name': str,
390
+ 'description': str,
391
+ 'scan_depth': (int, float),
392
+ 'token_budget': (int, float),
393
+ 'recursive_scanning': bool,
394
+ 'extensions': dict,
395
+ 'entries': list
396
+ }
397
+
398
+ for field, expected_type in optional_fields.items():
399
+ if field in book_data:
400
+ if not isinstance(book_data[field], expected_type):
401
+ validation_messages.append(f"Field 'character_book.{field}' must be of type '{expected_type}'.")
402
+ # 'entries' is required
403
+ if 'entries' not in book_data or not isinstance(book_data['entries'], list):
404
+ validation_messages.append("Field 'character_book.entries' is required and must be a list.")
405
+ return False, validation_messages
406
+
407
+ # Validate each entry in 'entries'
408
+ entries = book_data.get('entries', [])
409
+ entry_ids = set()
410
+ for idx, entry in enumerate(entries):
411
+ is_valid_entry, entry_messages = validate_character_book_entry(entry, idx, entry_ids)
412
+ if not is_valid_entry:
413
+ validation_messages.extend(entry_messages)
414
+
415
+ is_valid = len(validation_messages) == 0
416
+ return is_valid, validation_messages
417
+
418
+ def validate_character_book_entry(entry, idx, entry_ids):
419
+ """
420
+ Validate an entry in the 'character_book.entries' list.
421
+
422
+ Args:
423
+ entry (dict): The entry data.
424
+ idx (int): The index of the entry in the list.
425
+ entry_ids (set): A set of existing entry IDs for uniqueness checking.
426
+
427
+ Returns:
428
+ Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages.
429
+ """
430
+ validation_messages = []
431
+ required_fields = {
432
+ 'keys': list,
433
+ 'content': str,
434
+ 'extensions': dict,
435
+ 'enabled': bool,
436
+ 'insertion_order': (int, float)
437
+ }
438
+
439
+ for field, expected_type in required_fields.items():
440
+ if field not in entry:
441
+ validation_messages.append(f"Entry {idx}: Missing required field '{field}'.")
442
+ elif not isinstance(entry[field], expected_type):
443
+ validation_messages.append(f"Entry {idx}: Field '{field}' must be of type '{expected_type}'.")
444
+ elif field == 'content' and not entry[field].strip():
445
+ validation_messages.append(f"Entry {idx}: Field 'content' cannot be empty.")
446
+ elif field == 'keys' and not entry[field]:
447
+ validation_messages.append(f"Entry {idx}: Field 'keys' cannot be empty.")
448
+
449
+ # Optional fields
450
+ optional_fields = {
451
+ 'case_sensitive': bool,
452
+ 'name': str,
453
+ 'priority': (int, float),
454
+ 'id': (int, float),
455
+ 'comment': str,
456
+ 'selective': bool,
457
+ 'secondary_keys': list,
458
+ 'constant': bool,
459
+ 'position': str # Should be 'before_char' or 'after_char'
460
+ }
461
+
462
+ for field, expected_type in optional_fields.items():
463
+ if field in entry and not isinstance(entry[field], expected_type):
464
+ validation_messages.append(f"Entry {idx}: Field '{field}' must be of type '{expected_type}'.")
465
+
466
+ # Validate 'position' value if present
467
+ if 'position' in entry:
468
+ if entry['position'] not in ['before_char', 'after_char']:
469
+ validation_messages.append(f"Entry {idx}: Field 'position' must be 'before_char' or 'after_char'.")
470
+
471
+ # Validate 'secondary_keys' if 'selective' is True
472
+ if entry.get('selective', False):
473
+ if 'secondary_keys' not in entry or not isinstance(entry['secondary_keys'], list):
474
+ validation_messages.append(f"Entry {idx}: 'secondary_keys' must be a list when 'selective' is True.")
475
+ elif not entry['secondary_keys']:
476
+ validation_messages.append(f"Entry {idx}: 'secondary_keys' cannot be empty when 'selective' is True.")
477
+
478
+ # Validate 'keys' list elements
479
+ if 'keys' in entry and isinstance(entry['keys'], list):
480
+ for i, key in enumerate(entry['keys']):
481
+ if not isinstance(key, str) or not key.strip():
482
+ validation_messages.append(f"Entry {idx}: Element {i} in 'keys' must be a non-empty string.")
483
+
484
+ # Validate 'secondary_keys' list elements
485
+ if 'secondary_keys' in entry and isinstance(entry['secondary_keys'], list):
486
+ for i, key in enumerate(entry['secondary_keys']):
487
+ if not isinstance(key, str) or not key.strip():
488
+ validation_messages.append(f"Entry {idx}: Element {i} in 'secondary_keys' must be a non-empty string.")
489
+
490
+ # Validate 'id' uniqueness
491
+ if 'id' in entry:
492
+ entry_id = entry['id']
493
+ if entry_id in entry_ids:
494
+ validation_messages.append \
495
+ (f"Entry {idx}: Duplicate 'id' value '{entry_id}'. Each entry 'id' must be unique.")
496
+ else:
497
+ entry_ids.add(entry_id)
498
+
499
+ # Validate 'extensions' keys are namespaced
500
+ if 'extensions' in entry and isinstance(entry['extensions'], dict):
501
+ for key in entry['extensions'].keys():
502
+ if '/' not in key and '_' not in key:
503
+ validation_messages.append \
504
+ (f"Entry {idx}: Extension key '{key}' in 'extensions' should be namespaced to prevent conflicts.")
505
+
506
+ is_valid = len(validation_messages) == 0
507
+ return is_valid, validation_messages
508
+
509
+ def validate_v2_card(card_data):
510
+ """
511
+ Validate a character card according to the V2 specification.
512
+
513
+ Args:
514
+ card_data (dict): The parsed character card data.
515
+
516
+ Returns:
517
+ Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages.
518
+ """
519
+ validation_messages = []
520
+
521
+ # Check top-level fields
522
+ if 'spec' not in card_data:
523
+ validation_messages.append("Missing 'spec' field.")
524
+ elif card_data['spec'] != 'chara_card_v2':
525
+ validation_messages.append(f"Invalid 'spec' value: {card_data['spec']}. Expected 'chara_card_v2'.")
526
+
527
+ if 'spec_version' not in card_data:
528
+ validation_messages.append("Missing 'spec_version' field.")
529
+ else:
530
+ # Ensure 'spec_version' is '2.0' or higher
531
+ try:
532
+ spec_version = float(card_data['spec_version'])
533
+ if spec_version < 2.0:
534
+ validation_messages.append \
535
+ (f"'spec_version' must be '2.0' or higher. Found '{card_data['spec_version']}'.")
536
+ except ValueError:
537
+ validation_messages.append \
538
+ (f"Invalid 'spec_version' format: {card_data['spec_version']}. Must be a number as a string.")
539
+
540
+ if 'data' not in card_data:
541
+ validation_messages.append("Missing 'data' field.")
542
+ return False, validation_messages # Cannot proceed without 'data' field
543
+
544
+ data = card_data['data']
545
+
546
+ # Required fields in 'data'
547
+ required_fields = ['name', 'description', 'personality', 'scenario', 'first_mes', 'mes_example']
548
+ for field in required_fields:
549
+ if field not in data:
550
+ validation_messages.append(f"Missing required field in 'data': '{field}'.")
551
+ elif not isinstance(data[field], str):
552
+ validation_messages.append(f"Field '{field}' must be a string.")
553
+ elif not data[field].strip():
554
+ validation_messages.append(f"Field '{field}' cannot be empty.")
555
+
556
+ # Optional fields with expected types
557
+ optional_fields = {
558
+ 'creator_notes': str,
559
+ 'system_prompt': str,
560
+ 'post_history_instructions': str,
561
+ 'alternate_greetings': list,
562
+ 'tags': list,
563
+ 'creator': str,
564
+ 'character_version': str,
565
+ 'extensions': dict,
566
+ 'character_book': dict # If present, should be a dict
567
+ }
568
+
569
+ for field, expected_type in optional_fields.items():
570
+ if field in data:
571
+ if not isinstance(data[field], expected_type):
572
+ validation_messages.append(f"Field '{field}' must be of type '{expected_type.__name__}'.")
573
+ elif field == 'extensions':
574
+ # Validate that extensions keys are properly namespaced
575
+ for key in data[field].keys():
576
+ if '/' not in key and '_' not in key:
577
+ validation_messages.append \
578
+ (f"Extension key '{key}' in 'extensions' should be namespaced to prevent conflicts.")
579
+
580
+ # If 'alternate_greetings' is present, check that it's a list of non-empty strings
581
+ if 'alternate_greetings' in data and isinstance(data['alternate_greetings'], list):
582
+ for idx, greeting in enumerate(data['alternate_greetings']):
583
+ if not isinstance(greeting, str) or not greeting.strip():
584
+ validation_messages.append(f"Element {idx} in 'alternate_greetings' must be a non-empty string.")
585
+
586
+ # If 'tags' is present, check that it's a list of non-empty strings
587
+ if 'tags' in data and isinstance(data['tags'], list):
588
+ for idx, tag in enumerate(data['tags']):
589
+ if not isinstance(tag, str) or not tag.strip():
590
+ validation_messages.append(f"Element {idx} in 'tags' must be a non-empty string.")
591
+
592
+ # Validate 'extensions' field
593
+ if 'extensions' in data and not isinstance(data['extensions'], dict):
594
+ validation_messages.append("Field 'extensions' must be a dictionary.")
595
+
596
+ # Validate 'character_book' if present
597
+ if 'character_book' in data:
598
+ is_valid_book, book_messages = validate_character_book(data['character_book'])
599
+ if not is_valid_book:
600
+ validation_messages.extend(book_messages)
601
+
602
+ is_valid = len(validation_messages) == 0
603
+ return is_valid, validation_messages
604
+
605
+ #
606
+ # End of File
607
+ ####################################################################################################
App_Function_Libraries/Character_Chat/__init__.py ADDED
File without changes
App_Function_Libraries/Chat.py ADDED
@@ -0,0 +1,439 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Chat.py
2
+ # Chat functions for interacting with the LLMs as chatbots
3
+ import base64
4
+ # Imports
5
+ import json
6
+ import logging
7
+ import os
8
+ import re
9
+ import tempfile
10
+ import time
11
+ from datetime import datetime
12
+ from pathlib import Path
13
+ #
14
+ # External Imports
15
+ #
16
+ # Local Imports
17
+ from App_Function_Libraries.DB.DB_Manager import get_conversation_name, save_chat_history_to_database
18
+ from App_Function_Libraries.LLM_API_Calls import chat_with_openai, chat_with_anthropic, chat_with_cohere, \
19
+ chat_with_groq, chat_with_openrouter, chat_with_deepseek, chat_with_mistral, chat_with_huggingface
20
+ from App_Function_Libraries.LLM_API_Calls_Local import chat_with_aphrodite, chat_with_local_llm, chat_with_ollama, \
21
+ chat_with_kobold, chat_with_llama, chat_with_oobabooga, chat_with_tabbyapi, chat_with_vllm, chat_with_custom_openai
22
+ from App_Function_Libraries.DB.SQLite_DB import load_media_content
23
+ from App_Function_Libraries.Utils.Utils import generate_unique_filename, load_and_log_configs
24
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
25
+ #
26
+ ####################################################################################################
27
+ #
28
+ # Functions:
29
+
30
+ def chat_api_call(api_endpoint, api_key, input_data, prompt, temp, system_message=None):
31
+ log_counter("chat_api_call_attempt", labels={"api_endpoint": api_endpoint})
32
+ start_time = time.time()
33
+ if not api_key:
34
+ api_key = None
35
+ model = None
36
+ try:
37
+ logging.info(f"Debug - Chat API Call - API Endpoint: {api_endpoint}")
38
+ logging.info(f"Debug - Chat API Call - API Key: {api_key}")
39
+ logging.info(f"Debug - Chat chat_api_call - API Endpoint: {api_endpoint}")
40
+ if api_endpoint.lower() == 'openai':
41
+ response = chat_with_openai(api_key, input_data, prompt, temp, system_message)
42
+
43
+ elif api_endpoint.lower() == 'anthropic':
44
+ # Retrieve the model from config
45
+ loaded_config_data = load_and_log_configs()
46
+ model = loaded_config_data['models']['anthropic'] if loaded_config_data else None
47
+ response = chat_with_anthropic(
48
+ api_key=api_key,
49
+ input_data=input_data,
50
+ model=model,
51
+ custom_prompt_arg=prompt,
52
+ system_prompt=system_message
53
+ )
54
+
55
+ elif api_endpoint.lower() == "cohere":
56
+ response = chat_with_cohere(
57
+ api_key,
58
+ input_data,
59
+ model=model,
60
+ custom_prompt_arg=prompt,
61
+ system_prompt=system_message,
62
+ temp=temp
63
+ )
64
+
65
+ elif api_endpoint.lower() == "groq":
66
+ response = chat_with_groq(api_key, input_data, prompt, temp, system_message)
67
+
68
+ elif api_endpoint.lower() == "openrouter":
69
+ response = chat_with_openrouter(api_key, input_data, prompt, temp, system_message)
70
+
71
+ elif api_endpoint.lower() == "deepseek":
72
+ response = chat_with_deepseek(api_key, input_data, prompt, temp, system_message)
73
+
74
+ elif api_endpoint.lower() == "mistral":
75
+ response = chat_with_mistral(api_key, input_data, prompt, temp, system_message)
76
+
77
+ elif api_endpoint.lower() == "llama.cpp":
78
+ response = chat_with_llama(input_data, prompt, temp, None, api_key, system_message)
79
+ elif api_endpoint.lower() == "kobold":
80
+ response = chat_with_kobold(input_data, api_key, prompt, temp, system_message)
81
+
82
+ elif api_endpoint.lower() == "ooba":
83
+ response = chat_with_oobabooga(input_data, api_key, prompt, temp, system_message)
84
+
85
+ elif api_endpoint.lower() == "tabbyapi":
86
+ response = chat_with_tabbyapi(input_data, prompt, temp, system_message)
87
+
88
+ elif api_endpoint.lower() == "vllm":
89
+ response = chat_with_vllm(input_data, prompt, system_message)
90
+
91
+ elif api_endpoint.lower() == "local-llm":
92
+ response = chat_with_local_llm(input_data, prompt, temp, system_message)
93
+
94
+ elif api_endpoint.lower() == "huggingface":
95
+ response = chat_with_huggingface(api_key, input_data, prompt, temp) # , system_message)
96
+
97
+ elif api_endpoint.lower() == "ollama":
98
+ response = chat_with_ollama(input_data, prompt, None, api_key, temp, system_message)
99
+
100
+ elif api_endpoint.lower() == "aphrodite":
101
+ response = chat_with_aphrodite(input_data, prompt, temp, system_message)
102
+
103
+ elif api_endpoint.lower() == "custom-openai-api":
104
+ response = chat_with_custom_openai(api_key, input_data, prompt, temp, system_message)
105
+
106
+ else:
107
+ raise ValueError(f"Unsupported API endpoint: {api_endpoint}")
108
+
109
+ call_duration = time.time() - start_time
110
+ log_histogram("chat_api_call_duration", call_duration, labels={"api_endpoint": api_endpoint})
111
+ log_counter("chat_api_call_success", labels={"api_endpoint": api_endpoint})
112
+ return response
113
+
114
+ except Exception as e:
115
+ log_counter("chat_api_call_error", labels={"api_endpoint": api_endpoint, "error": str(e)})
116
+ logging.error(f"Error in chat function: {str(e)}")
117
+ return f"An error occurred: {str(e)}"
118
+
119
+
120
+ def chat(message, history, media_content, selected_parts, api_endpoint, api_key, prompt, temperature,
121
+ system_message=None):
122
+ log_counter("chat_attempt", labels={"api_endpoint": api_endpoint})
123
+ start_time = time.time()
124
+ try:
125
+ logging.info(f"Debug - Chat Function - Message: {message}")
126
+ logging.info(f"Debug - Chat Function - Media Content: {media_content}")
127
+ logging.info(f"Debug - Chat Function - Selected Parts: {selected_parts}")
128
+ logging.info(f"Debug - Chat Function - API Endpoint: {api_endpoint}")
129
+ # logging.info(f"Debug - Chat Function - Prompt: {prompt}")
130
+
131
+ # Ensure selected_parts is a list
132
+ if not isinstance(selected_parts, (list, tuple)):
133
+ selected_parts = [selected_parts] if selected_parts else []
134
+
135
+ # logging.debug(f"Debug - Chat Function - Selected Parts (after check): {selected_parts}")
136
+
137
+ # Combine the selected parts of the media content
138
+ combined_content = "\n\n".join(
139
+ [f"{part.capitalize()}: {media_content.get(part, '')}" for part in selected_parts if part in media_content])
140
+ # Print first 500 chars
141
+ # logging.debug(f"Debug - Chat Function - Combined Content: {combined_content[:500]}...")
142
+
143
+ # Prepare the input for the API
144
+ input_data = f"{combined_content}\n\n" if combined_content else ""
145
+ for old_message, old_response in history:
146
+ input_data += f"{old_message}\nAssistant: {old_response}\n\n"
147
+ input_data += f"{message}\n"
148
+
149
+ if system_message:
150
+ print(f"System message: {system_message}")
151
+ logging.debug(f"Debug - Chat Function - System Message: {system_message}")
152
+ temperature = float(temperature) if temperature else 0.7
153
+ temp = temperature
154
+
155
+ logging.debug(f"Debug - Chat Function - Temperature: {temperature}")
156
+ logging.debug(f"Debug - Chat Function - API Key: {api_key[:10]}")
157
+ logging.debug(f"Debug - Chat Function - Prompt: {prompt}")
158
+
159
+ # Use the existing API request code based on the selected endpoint
160
+ response = chat_api_call(api_endpoint, api_key, input_data, prompt, temp, system_message)
161
+
162
+ chat_duration = time.time() - start_time
163
+ log_histogram("chat_duration", chat_duration, labels={"api_endpoint": api_endpoint})
164
+ log_counter("chat_success", labels={"api_endpoint": api_endpoint})
165
+ return response
166
+ except Exception as e:
167
+ log_counter("chat_error", labels={"api_endpoint": api_endpoint, "error": str(e)})
168
+ logging.error(f"Error in chat function: {str(e)}")
169
+ return f"An error occurred: {str(e)}"
170
+
171
+
172
+ def save_chat_history_to_db_wrapper(chatbot, conversation_id, media_content, media_name=None):
173
+ log_counter("save_chat_history_to_db_attempt")
174
+ start_time = time.time()
175
+ logging.info(f"Attempting to save chat history. Media content type: {type(media_content)}")
176
+ try:
177
+ # Extract the media_id and media_name from the media_content
178
+ media_id = None
179
+ if isinstance(media_content, dict):
180
+ media_id = None
181
+ logging.debug(f"Media content keys: {media_content.keys()}")
182
+ if 'content' in media_content:
183
+ try:
184
+ content = media_content['content']
185
+ if isinstance(content, str):
186
+ content_json = json.loads(content)
187
+ elif isinstance(content, dict):
188
+ content_json = content
189
+ else:
190
+ raise ValueError(f"Unexpected content type: {type(content)}")
191
+
192
+ # Use the webpage_url as the media_id
193
+ media_id = content_json.get('webpage_url')
194
+ # Use the title as the media_name
195
+ media_name = content_json.get('title')
196
+
197
+ logging.info(f"Extracted media_id: {media_id}, media_name: {media_name}")
198
+ except json.JSONDecodeError:
199
+ logging.error("Failed to decode JSON from media_content['content']")
200
+ except Exception as e:
201
+ logging.error(f"Error processing media_content: {str(e)}")
202
+ else:
203
+ logging.warning("'content' key not found in media_content")
204
+ else:
205
+ logging.warning(f"media_content is not a dictionary. Type: {type(media_content)}")
206
+
207
+ if media_id is None:
208
+ # If we couldn't find a media_id, we'll use a placeholder
209
+ media_id = "unknown_media"
210
+ logging.warning(f"Unable to extract media_id from media_content. Using placeholder: {media_id}")
211
+
212
+ if media_name is None:
213
+ media_name = "Unnamed Media"
214
+ logging.warning(f"Unable to extract media_name from media_content. Using placeholder: {media_name}")
215
+
216
+ # Generate a unique conversation name using media_id and current timestamp
217
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
218
+ conversation_name = f"{media_name}_{timestamp}"
219
+
220
+ new_conversation_id = save_chat_history_to_database(chatbot, conversation_id, media_id, media_name,
221
+ conversation_name)
222
+ save_duration = time.time() - start_time
223
+ log_histogram("save_chat_history_to_db_duration", save_duration)
224
+ log_counter("save_chat_history_to_db_success")
225
+ return new_conversation_id, f"Chat history saved successfully as {conversation_name}!"
226
+ except Exception as e:
227
+ log_counter("save_chat_history_to_db_error", labels={"error": str(e)})
228
+ error_message = f"Failed to save chat history: {str(e)}"
229
+ logging.error(error_message, exc_info=True)
230
+ return conversation_id, error_message
231
+
232
+
233
+ def save_chat_history(history, conversation_id, media_content):
234
+ log_counter("save_chat_history_attempt")
235
+ start_time = time.time()
236
+ try:
237
+ content, conversation_name = generate_chat_history_content(history, conversation_id, media_content)
238
+
239
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
240
+ safe_conversation_name = re.sub(r'[^a-zA-Z0-9_-]', '_', conversation_name)
241
+ base_filename = f"{safe_conversation_name}_{timestamp}.json"
242
+
243
+ # Create a temporary file
244
+ with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as temp_file:
245
+ temp_file.write(content)
246
+ temp_file_path = temp_file.name
247
+
248
+ # Generate a unique filename
249
+ unique_filename = generate_unique_filename(os.path.dirname(temp_file_path), base_filename)
250
+ final_path = os.path.join(os.path.dirname(temp_file_path), unique_filename)
251
+
252
+ # Rename the temporary file to the unique filename
253
+ os.rename(temp_file_path, final_path)
254
+
255
+ save_duration = time.time() - start_time
256
+ log_histogram("save_chat_history_duration", save_duration)
257
+ log_counter("save_chat_history_success")
258
+ return final_path
259
+ except Exception as e:
260
+ log_counter("save_chat_history_error", labels={"error": str(e)})
261
+ logging.error(f"Error saving chat history: {str(e)}")
262
+ return None
263
+
264
+
265
+ def generate_chat_history_content(history, conversation_id, media_content):
266
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
267
+
268
+ conversation_name = get_conversation_name(conversation_id)
269
+
270
+ if not conversation_name:
271
+ media_name = extract_media_name(media_content)
272
+ if media_name:
273
+ conversation_name = f"{media_name}-chat"
274
+ else:
275
+ conversation_name = f"chat-{timestamp}" # Fallback name
276
+
277
+ chat_data = {
278
+ "conversation_id": conversation_id,
279
+ "conversation_name": conversation_name,
280
+ "timestamp": timestamp,
281
+ "history": [
282
+ {
283
+ "role": "user" if i % 2 == 0 else "bot",
284
+ "content": msg[0] if isinstance(msg, tuple) else msg
285
+ }
286
+ for i, msg in enumerate(history)
287
+ ]
288
+ }
289
+
290
+ return json.dumps(chat_data, indent=2), conversation_name
291
+
292
+
293
+ def extract_media_name(media_content):
294
+ if isinstance(media_content, dict):
295
+ content = media_content.get('content', {})
296
+ if isinstance(content, str):
297
+ try:
298
+ content = json.loads(content)
299
+ except json.JSONDecodeError:
300
+ logging.warning("Failed to parse media_content JSON string")
301
+ return None
302
+
303
+ # Try to extract title from the content
304
+ if isinstance(content, dict):
305
+ return content.get('title') or content.get('name')
306
+
307
+ logging.warning(f"Unexpected media_content format: {type(media_content)}")
308
+ return None
309
+
310
+
311
+ def update_chat_content(selected_item, use_content, use_summary, use_prompt, item_mapping):
312
+ log_counter("update_chat_content_attempt")
313
+ start_time = time.time()
314
+ logging.debug(f"Debug - Update Chat Content - Selected Item: {selected_item}\n")
315
+ logging.debug(f"Debug - Update Chat Content - Use Content: {use_content}\n\n\n\n")
316
+ logging.debug(f"Debug - Update Chat Content - Use Summary: {use_summary}\n\n")
317
+ logging.debug(f"Debug - Update Chat Content - Use Prompt: {use_prompt}\n\n")
318
+ logging.debug(f"Debug - Update Chat Content - Item Mapping: {item_mapping}\n\n")
319
+
320
+ if selected_item and selected_item in item_mapping:
321
+ media_id = item_mapping[selected_item]
322
+ content = load_media_content(media_id)
323
+ selected_parts = []
324
+ if use_content and "content" in content:
325
+ selected_parts.append("content")
326
+ if use_summary and "summary" in content:
327
+ selected_parts.append("summary")
328
+ if use_prompt and "prompt" in content:
329
+ selected_parts.append("prompt")
330
+
331
+ # Modified debug print
332
+ if isinstance(content, dict):
333
+ print(f"Debug - Update Chat Content - Content keys: {list(content.keys())}")
334
+ for key, value in content.items():
335
+ print(f"Debug - Update Chat Content - {key} (first 500 char): {str(value)[:500]}\n\n\n\n")
336
+ else:
337
+ print(f"Debug - Update Chat Content - Content(first 500 char): {str(content)[:500]}\n\n\n\n")
338
+
339
+ print(f"Debug - Update Chat Content - Selected Parts: {selected_parts}")
340
+ update_duration = time.time() - start_time
341
+ log_histogram("update_chat_content_duration", update_duration)
342
+ log_counter("update_chat_content_success")
343
+ return content, selected_parts
344
+ else:
345
+ log_counter("update_chat_content_error", labels={"error": str("No item selected or item not in mapping")})
346
+ print(f"Debug - Update Chat Content - No item selected or item not in mapping")
347
+ return {}, []
348
+
349
+ #
350
+ # End of Chat functions
351
+ #######################################################################################################################
352
+
353
+
354
+ #######################################################################################################################
355
+ #
356
+ # Character Card Functions
357
+
358
+ CHARACTERS_FILE = Path('.', 'Helper_Scripts', 'Character_Cards', 'Characters.json')
359
+
360
+
361
+ def save_character(character_data):
362
+ log_counter("save_character_attempt")
363
+ start_time = time.time()
364
+ characters_file = os.path.join(os.path.dirname(__file__), '..', 'Helper_Scripts', 'Character_Cards', 'Characters.json')
365
+ characters_dir = os.path.dirname(characters_file)
366
+
367
+ try:
368
+ if os.path.exists(characters_file):
369
+ with open(characters_file, 'r') as f:
370
+ characters = json.load(f)
371
+ else:
372
+ characters = {}
373
+
374
+ char_name = character_data['name']
375
+
376
+ # Save the image separately if it exists
377
+ if 'image' in character_data:
378
+ img_data = base64.b64decode(character_data['image'])
379
+ img_filename = f"{char_name.replace(' ', '_')}.png"
380
+ img_path = os.path.join(characters_dir, img_filename)
381
+ with open(img_path, 'wb') as f:
382
+ f.write(img_data)
383
+ character_data['image_path'] = os.path.abspath(img_path)
384
+ del character_data['image'] # Remove the base64 image data from the JSON
385
+
386
+ characters[char_name] = character_data
387
+
388
+ with open(characters_file, 'w') as f:
389
+ json.dump(characters, f, indent=2)
390
+
391
+ save_duration = time.time() - start_time
392
+ log_histogram("save_character_duration", save_duration)
393
+ log_counter("save_character_success")
394
+ logging.info(f"Character '{char_name}' saved successfully.")
395
+ except Exception as e:
396
+ log_counter("save_character_error", labels={"error": str(e)})
397
+ logging.error(f"Error saving character: {str(e)}")
398
+
399
+
400
+ def load_characters():
401
+ log_counter("load_characters_attempt")
402
+ start_time = time.time()
403
+ try:
404
+ characters_file = os.path.join(os.path.dirname(__file__), '..', 'Helper_Scripts', 'Character_Cards', 'Characters.json')
405
+ if os.path.exists(characters_file):
406
+ with open(characters_file, 'r') as f:
407
+ characters = json.load(f)
408
+ logging.debug(f"Loaded {len(characters)} characters from {characters_file}")
409
+ load_duration = time.time() - start_time
410
+ log_histogram("load_characters_duration", load_duration)
411
+ log_counter("load_characters_success", labels={"character_count": len(characters)})
412
+ return characters
413
+ else:
414
+ logging.warning(f"Characters file not found: {characters_file}")
415
+ return {}
416
+ except Exception as e:
417
+ log_counter("load_characters_error", labels={"error": str(e)})
418
+ return {}
419
+
420
+
421
+
422
+ def get_character_names():
423
+ log_counter("get_character_names_attempt")
424
+ start_time = time.time()
425
+ try:
426
+ characters = load_characters()
427
+ names = list(characters.keys())
428
+ get_names_duration = time.time() - start_time
429
+ log_histogram("get_character_names_duration", get_names_duration)
430
+ log_counter("get_character_names_success", labels={"name_count": len(names)})
431
+ return names
432
+ except Exception as e:
433
+ log_counter("get_character_names_error", labels={"error": str(e)})
434
+ logging.error(f"Error getting character names: {str(e)}")
435
+ return []
436
+
437
+ #
438
+ # End of Chat.py
439
+ ##########################################################################################################################
App_Function_Libraries/Chunk_Lib.py ADDED
@@ -0,0 +1,1051 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Chunk_Lib.py
2
+ #########################################
3
+ # Chunking Library
4
+ # This library is used to perform chunking of input files.
5
+ # Currently, uses naive approaches. Nothing fancy.
6
+ #
7
+ ####
8
+ # Import necessary libraries
9
+ import hashlib
10
+ import json
11
+ import logging
12
+ import re
13
+ from typing import Any, Dict, List, Optional, Tuple
14
+ #
15
+ # Import 3rd party
16
+ from openai import OpenAI
17
+ from tqdm import tqdm
18
+ from langdetect import detect
19
+ from transformers import GPT2Tokenizer
20
+ import nltk
21
+ from nltk.tokenize import sent_tokenize, word_tokenize
22
+ from sklearn.feature_extraction.text import TfidfVectorizer
23
+ from sklearn.metrics.pairwise import cosine_similarity
24
+ #
25
+ # Import Local
26
+ from App_Function_Libraries.Tokenization_Methods_Lib import openai_tokenize
27
+ from App_Function_Libraries.Utils.Utils import load_comprehensive_config
28
+ #
29
+ #######################################################################################################################
30
+ # Config Settings
31
+ #
32
+ #
33
+ # FIXME - Make sure it only downloads if it already exists, and does a check first.
34
+ # Ensure NLTK data is downloaded
35
+ def ensure_nltk_data():
36
+ try:
37
+ nltk.data.find('tokenizers/punkt')
38
+ except LookupError:
39
+ nltk.download('punkt')
40
+ ensure_nltk_data()
41
+
42
+ #
43
+ # Load GPT2 tokenizer
44
+ tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
45
+ #
46
+ # Load configuration
47
+ config = load_comprehensive_config()
48
+ # Embedding Chunking options
49
+ chunk_options = {
50
+ 'method': config.get('Chunking', 'method', fallback='words'),
51
+ 'max_size': config.getint('Chunking', 'max_size', fallback=400),
52
+ 'overlap': config.getint('Chunking', 'overlap', fallback=200),
53
+ 'adaptive': config.getboolean('Chunking', 'adaptive', fallback=False),
54
+ 'multi_level': config.getboolean('Chunking', 'multi_level', fallback=False),
55
+ 'language': config.get('Chunking', 'language', fallback='english')
56
+ }
57
+
58
+ openai_api_key = config.get('API', 'openai_api_key')
59
+ #
60
+ # End of settings
61
+ #######################################################################################################################
62
+ #
63
+ # Functions:
64
+
65
+ # Create a chunking class for refactoring FIXME
66
+ # class Chunker:
67
+ # def __init__(self, tokenizer: GPT2Tokenizer):
68
+ # self.tokenizer = tokenizer
69
+ #
70
+ # def detect_language(self, text: str) -> str:
71
+ # try:
72
+ # return detect(text)
73
+ # except:
74
+ # return 'en'
75
+ #
76
+ # def chunk_text(self, text: str, method: str, max_size: int, overlap: int, language: str = None) -> List[str]:
77
+ # if language is None:
78
+ # language = self.detect_language(text)
79
+ #
80
+ # if method == 'words':
81
+ # return self.chunk_text_by_words(text, max_size, overlap, language)
82
+ # elif method == 'sentences':
83
+ # return self.chunk_text_by_sentences(text, max_size, overlap, language)
84
+ # elif method == 'paragraphs':
85
+ # return self.chunk_text_by_paragraphs(text, max_size, overlap)
86
+ # elif method == 'tokens':
87
+ # return self.chunk_text_by_tokens(text, max_size, overlap, language)
88
+ # elif method == 'semantic':
89
+ # return self.semantic_chunking(text, max_size)
90
+ # else:
91
+ # return [text]
92
+
93
+ def detect_language(text: str) -> str:
94
+ try:
95
+ return detect(text)
96
+ except:
97
+ # Default to English if detection fails
98
+ return 'en'
99
+
100
+
101
+ def load_document(file_path: str) -> str:
102
+ with open(file_path, 'r', encoding='utf-8') as file:
103
+ text = file.read()
104
+ return re.sub(r'\s+', ' ', text).strip()
105
+
106
+
107
+ def improved_chunking_process(text: str, chunk_options: Dict[str, Any] = None) -> List[Dict[str, Any]]:
108
+ logging.debug("Improved chunking process started...")
109
+
110
+ # Extract JSON metadata if present
111
+ json_content = {}
112
+ try:
113
+ json_end = text.index("}\n") + 1
114
+ json_content = json.loads(text[:json_end])
115
+ text = text[json_end:].strip()
116
+ logging.debug(f"Extracted JSON metadata: {json_content}")
117
+ except (ValueError, json.JSONDecodeError):
118
+ logging.debug("No JSON metadata found at the beginning of the text")
119
+
120
+ # Extract any additional header text
121
+ header_match = re.match(r"(This text was transcribed using.*?)\n\n", text, re.DOTALL)
122
+ header_text = ""
123
+ if header_match:
124
+ header_text = header_match.group(1)
125
+ text = text[len(header_text):].strip()
126
+ logging.debug(f"Extracted header text: {header_text}")
127
+
128
+ options = chunk_options.copy() if chunk_options else {}
129
+ if chunk_options:
130
+ options.update(chunk_options)
131
+
132
+ chunk_method = options.get('method', 'words')
133
+ max_size = options.get('max_size', 2000)
134
+ overlap = options.get('overlap', 0)
135
+ language = options.get('language', None)
136
+
137
+ if language is None:
138
+ language = detect_language(text)
139
+
140
+ if chunk_method == 'json':
141
+ chunks = chunk_text_by_json(text, max_size=max_size, overlap=overlap)
142
+ else:
143
+ chunks = chunk_text(text, chunk_method, max_size, overlap, language)
144
+
145
+ chunks_with_metadata = []
146
+ total_chunks = len(chunks)
147
+ for i, chunk in enumerate(chunks):
148
+ metadata = {
149
+ 'chunk_index': i + 1,
150
+ 'total_chunks': total_chunks,
151
+ 'chunk_method': chunk_method,
152
+ 'max_size': max_size,
153
+ 'overlap': overlap,
154
+ 'language': language,
155
+ 'relative_position': (i + 1) / total_chunks
156
+ }
157
+ metadata.update(json_content) # Add the extracted JSON content to metadata
158
+ metadata['header_text'] = header_text # Add the header text to metadata
159
+
160
+ if chunk_method == 'json':
161
+ chunk_text_content = json.dumps(chunk['json'], ensure_ascii=False)
162
+ else:
163
+ chunk_text_content = chunk
164
+
165
+ chunks_with_metadata.append({
166
+ 'text': chunk_text_content,
167
+ 'metadata': metadata
168
+ })
169
+
170
+ return chunks_with_metadata
171
+
172
+
173
+ def multi_level_chunking(text: str, method: str, max_size: int, overlap: int, language: str) -> List[str]:
174
+ logging.debug("Multi-level chunking process started...")
175
+ # First level: chunk by paragraphs
176
+ paragraphs = chunk_text_by_paragraphs(text, max_size * 2, overlap)
177
+
178
+ # Second level: chunk each paragraph further
179
+ chunks = []
180
+ for para in paragraphs:
181
+ if method == 'words':
182
+ chunks.extend(chunk_text_by_words(para, max_words=max_size, overlap=overlap, language=language))
183
+ elif method == 'sentences':
184
+ chunks.extend(chunk_text_by_sentences(para, max_sentences=max_size, overlap=overlap, language=language))
185
+ else:
186
+ chunks.append(para)
187
+
188
+ return chunks
189
+
190
+
191
+ # FIXME - ensure language detection occurs in each chunk function
192
+ def chunk_text(text: str, method: str, max_size: int, overlap: int, language: str = None) -> List[str]:
193
+ if method == 'words':
194
+ logging.debug("Chunking by words...")
195
+ return chunk_text_by_words(text, max_words=max_size, overlap=overlap, language=language)
196
+ elif method == 'sentences':
197
+ logging.debug("Chunking by sentences...")
198
+ return chunk_text_by_sentences(text, max_sentences=max_size, overlap=overlap, language=language)
199
+ elif method == 'paragraphs':
200
+ logging.debug("Chunking by paragraphs...")
201
+ return chunk_text_by_paragraphs(text, max_paragraphs=max_size, overlap=overlap)
202
+ elif method == 'tokens':
203
+ logging.debug("Chunking by tokens...")
204
+ return chunk_text_by_tokens(text, max_tokens=max_size, overlap=overlap)
205
+ elif method == 'semantic':
206
+ logging.debug("Chunking by semantic similarity...")
207
+ return semantic_chunking(text, max_chunk_size=max_size)
208
+ else:
209
+ logging.warning(f"Unknown chunking method '{method}'. Returning full text as a single chunk.")
210
+ return [text]
211
+
212
+ def determine_chunk_position(relative_position: float) -> str:
213
+ if relative_position < 0.33:
214
+ return "This chunk is from the beginning of the document"
215
+ elif relative_position < 0.66:
216
+ return "This chunk is from the middle of the document"
217
+ else:
218
+ return "This chunk is from the end of the document"
219
+
220
+
221
+ def chunk_text_by_words(text: str, max_words: int = 300, overlap: int = 0, language: str = None) -> List[str]:
222
+ logging.debug("chunk_text_by_words...")
223
+ if language is None:
224
+ language = detect_language(text)
225
+
226
+ if language.startswith('zh'): # Chinese
227
+ import jieba
228
+ words = list(jieba.cut(text))
229
+ elif language == 'ja': # Japanese
230
+ import fugashi
231
+ tagger = fugashi.Tagger()
232
+ words = [word.surface for word in tagger(text)]
233
+ else: # Default to simple splitting for other languages
234
+ words = text.split()
235
+
236
+ chunks = []
237
+ for i in range(0, len(words), max_words - overlap):
238
+ chunk = ' '.join(words[i:i + max_words])
239
+ chunks.append(chunk)
240
+ return post_process_chunks(chunks)
241
+
242
+
243
+ def chunk_text_by_sentences(text: str, max_sentences: int = 10, overlap: int = 0, language: str = None) -> List[str]:
244
+ logging.debug("chunk_text_by_sentences...")
245
+ if language is None:
246
+ language = detect_language(text)
247
+
248
+ if language.startswith('zh'): # Chinese
249
+ import jieba
250
+ # Use jieba to perform sentence segmentation
251
+ # jieba does not support sentence segmentation out of the box
252
+ # Use punctuation as delimiters
253
+ sentences = re.split(r'[。!?;]', text)
254
+ sentences = [s.strip() for s in sentences if s.strip()]
255
+ elif language == 'ja': # Japanese
256
+ import fugashi
257
+ tagger = fugashi.Tagger()
258
+ # Simple sentence segmentation based on punctuation
259
+ sentences = re.split(r'[。!?]', text)
260
+ sentences = [s.strip() for s in sentences if s.strip()]
261
+ else: # Default to NLTK for other languages
262
+ try:
263
+ sentences = sent_tokenize(text, language=language)
264
+ except LookupError:
265
+ logging.warning(f"Punkt tokenizer not found for language '{language}'. Using default 'english'.")
266
+ sentences = sent_tokenize(text, language='english')
267
+
268
+ chunks = []
269
+ previous_overlap = []
270
+
271
+ for i in range(0, len(sentences), max_sentences - overlap):
272
+ current_sentences = sentences[i:i + max_sentences]
273
+ if overlap > 0 and previous_overlap:
274
+ current_sentences = previous_overlap + current_sentences
275
+ chunk = ' '.join(current_sentences)
276
+ chunks.append(chunk)
277
+ previous_overlap = sentences[i + max_sentences - overlap:i + max_sentences] if overlap > 0 else []
278
+
279
+ return post_process_chunks(chunks)
280
+
281
+
282
+ def chunk_text_by_paragraphs(text: str, max_paragraphs: int = 5, overlap: int = 0) -> List[str]:
283
+ logging.debug("chunk_text_by_paragraphs...")
284
+ paragraphs = re.split(r'\n\s*\n', text)
285
+ chunks = []
286
+ for i in range(0, len(paragraphs), max_paragraphs - overlap):
287
+ chunk = '\n\n'.join(paragraphs[i:i + max_paragraphs])
288
+ chunks.append(chunk)
289
+ return post_process_chunks(chunks)
290
+
291
+
292
+ def chunk_text_by_tokens(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]:
293
+ logging.debug("chunk_text_by_tokens...")
294
+ # This is a simplified token-based chunking. For more accurate tokenization,
295
+ # consider using a proper tokenizer like GPT-2 TokenizerFast
296
+ words = text.split()
297
+ chunks = []
298
+ current_chunk = []
299
+ current_token_count = 0
300
+
301
+ for word in words:
302
+ word_token_count = len(word) // 4 + 1 # Rough estimate of token count
303
+ if current_token_count + word_token_count > max_tokens and current_chunk:
304
+ chunks.append(' '.join(current_chunk))
305
+ current_chunk = current_chunk[-overlap:] if overlap > 0 else []
306
+ current_token_count = sum(len(w) // 4 + 1 for w in current_chunk)
307
+
308
+ current_chunk.append(word)
309
+ current_token_count += word_token_count
310
+
311
+ if current_chunk:
312
+ chunks.append(' '.join(current_chunk))
313
+
314
+ return post_process_chunks(chunks)
315
+ # def chunk_text_by_tokens(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]:
316
+ # logging.debug("chunk_text_by_tokens...")
317
+ # # Use GPT2 tokenizer for tokenization
318
+ # tokens = tokenizer.encode(text)
319
+ # chunks = []
320
+ # for i in range(0, len(tokens), max_tokens - overlap):
321
+ # chunk_tokens = tokens[i:i + max_tokens]
322
+ # chunk = tokenizer.decode(chunk_tokens)
323
+ # chunks.append(chunk)
324
+ # return post_process_chunks(chunks)
325
+
326
+
327
+ def post_process_chunks(chunks: List[str]) -> List[str]:
328
+ return [chunk.strip() for chunk in chunks if chunk.strip()]
329
+
330
+
331
+ # FIXME - F
332
+ def get_chunk_metadata(chunk: str, full_text: str, chunk_type: str = "generic",
333
+ chapter_number: Optional[int] = None,
334
+ chapter_pattern: Optional[str] = None,
335
+ language: str = None) -> Dict[str, Any]:
336
+ """
337
+ Generate metadata for a chunk based on its position in the full text.
338
+ """
339
+ chunk_length = len(chunk)
340
+ start_index = full_text.find(chunk)
341
+ end_index = start_index + chunk_length if start_index != -1 else None
342
+
343
+ # Calculate a hash for the chunk
344
+ chunk_hash = hashlib.md5(chunk.encode()).hexdigest()
345
+
346
+ metadata = {
347
+ 'start_index': start_index,
348
+ 'end_index': end_index,
349
+ 'word_count': len(chunk.split()),
350
+ 'char_count': chunk_length,
351
+ 'chunk_type': chunk_type,
352
+ 'language': language,
353
+ 'chunk_hash': chunk_hash,
354
+ 'relative_position': start_index / len(full_text) if len(full_text) > 0 and start_index != -1 else 0
355
+ }
356
+
357
+ if chunk_type == "chapter":
358
+ metadata['chapter_number'] = chapter_number
359
+ metadata['chapter_pattern'] = chapter_pattern
360
+
361
+ return metadata
362
+
363
+
364
+ def process_document_with_metadata(text: str, chunk_options: Dict[str, Any],
365
+ document_metadata: Dict[str, Any]) -> Dict[str, Any]:
366
+ chunks = improved_chunking_process(text, chunk_options)
367
+
368
+ return {
369
+ 'document_metadata': document_metadata,
370
+ 'chunks': chunks
371
+ }
372
+
373
+
374
+ # Hybrid approach, chunk each sentence while ensuring total token size does not exceed a maximum number
375
+ def chunk_text_hybrid(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]:
376
+ logging.debug("chunk_text_hybrid...")
377
+ sentences = sent_tokenize(text)
378
+ chunks = []
379
+ current_chunk = []
380
+ current_length = 0
381
+
382
+ for sentence in sentences:
383
+ tokens = tokenizer.encode(sentence)
384
+ if current_length + len(tokens) > max_tokens and current_chunk:
385
+ chunks.append(' '.join(current_chunk))
386
+ # Handle overlap
387
+ if overlap > 0:
388
+ overlap_tokens = tokenizer.encode(' '.join(current_chunk[-overlap:]))
389
+ current_chunk = current_chunk[-overlap:]
390
+ current_length = len(overlap_tokens)
391
+ else:
392
+ current_chunk = []
393
+ current_length = 0
394
+
395
+ current_chunk.append(sentence)
396
+ current_length += len(tokens)
397
+
398
+ if current_chunk:
399
+ chunks.append(' '.join(current_chunk))
400
+
401
+ return post_process_chunks(chunks)
402
+
403
+
404
+ # Thanks openai
405
+ def chunk_on_delimiter(input_string: str,
406
+ max_tokens: int,
407
+ delimiter: str) -> List[str]:
408
+ logging.debug("chunk_on_delimiter...")
409
+ chunks = input_string.split(delimiter)
410
+ combined_chunks, _, dropped_chunk_count = combine_chunks_with_no_minimum(
411
+ chunks, max_tokens, chunk_delimiter=delimiter, add_ellipsis_for_overflow=True)
412
+ if dropped_chunk_count > 0:
413
+ logging.warning(f"Warning: {dropped_chunk_count} chunks were dropped due to exceeding the token limit.")
414
+ combined_chunks = [f"{chunk}{delimiter}" for chunk in combined_chunks]
415
+ return combined_chunks
416
+
417
+
418
+
419
+
420
+ # FIXME
421
+ def recursive_summarize_chunks(chunks: List[str], summarize_func, custom_prompt: Optional[str] = None,
422
+ temp: Optional[float] = None, system_prompt: Optional[str] = None) -> List[str]:
423
+ logging.debug("recursive_summarize_chunks...")
424
+ summarized_chunks = []
425
+ current_summary = ""
426
+
427
+ logging.debug(f"Summarizing {len(chunks)} chunks recursively...")
428
+ logging.debug(f"Temperature is set to {temp}")
429
+ for i, chunk in enumerate(chunks):
430
+ if i == 0:
431
+ current_summary = summarize_func(chunk, custom_prompt, temp, system_prompt)
432
+ else:
433
+ combined_text = current_summary + "\n\n" + chunk
434
+ current_summary = summarize_func(combined_text, custom_prompt, temp, system_prompt)
435
+
436
+ summarized_chunks.append(current_summary)
437
+
438
+ return summarized_chunks
439
+
440
+
441
+ # Sample text for testing
442
+ sample_text = """
443
+ Natural language processing (NLP) is a subfield of linguistics, computer science, and artificial intelligence
444
+ concerned with the interactions between computers and human language, in particular how to program computers
445
+ to process and analyze large amounts of natural language data. The result is a computer capable of "understanding"
446
+ the contents of documents, including the contextual nuances of the language within them. The technology can then
447
+ accurately extract information and insights contained in the documents as well as categorize and organize the documents themselves.
448
+
449
+ Challenges in natural language processing frequently involve speech recognition, natural language understanding,
450
+ and natural language generation.
451
+
452
+ Natural language processing has its roots in the 1950s. Already in 1950, Alan Turing published an article titled
453
+ "Computing Machinery and Intelligence" which proposed what is now called the Turing test as a criterion of intelligence.
454
+ """
455
+
456
+ # Example usage of different chunking methods
457
+ # print("Chunking by words:")
458
+ # print(chunk_text_by_words(sample_text, max_words=50))
459
+ #
460
+ # print("\nChunking by sentences:")
461
+ # print(chunk_text_by_sentences(sample_text, max_sentences=2))
462
+ #
463
+ # print("\nChunking by paragraphs:")
464
+ # print(chunk_text_by_paragraphs(sample_text, max_paragraphs=1))
465
+ #
466
+ # print("\nChunking by tokens:")
467
+ # print(chunk_text_by_tokens(sample_text, max_tokens=50))
468
+ #
469
+ # print("\nHybrid chunking:")
470
+ # print(chunk_text_hybrid(sample_text, max_tokens=50))
471
+
472
+
473
+
474
+ #######################################################################################################################
475
+ #
476
+ # Experimental Semantic Chunking
477
+ #
478
+
479
+ # Chunk text into segments based on semantic similarity
480
+ def count_units(text: str, unit: str = 'words') -> int:
481
+ if unit == 'words':
482
+ return len(text.split())
483
+ elif unit == 'tokens':
484
+ return len(tokenizer.encode(text))
485
+ elif unit == 'characters':
486
+ return len(text)
487
+ else:
488
+ raise ValueError("Invalid unit. Choose 'words', 'tokens', or 'characters'.")
489
+
490
+
491
+
492
+ def semantic_chunking(text: str, max_chunk_size: int = 2000, unit: str = 'words') -> List[str]:
493
+ logging.debug("semantic_chunking...")
494
+ sentences = sent_tokenize(text)
495
+ vectorizer = TfidfVectorizer()
496
+ sentence_vectors = vectorizer.fit_transform(sentences)
497
+
498
+ chunks = []
499
+ current_chunk = []
500
+ current_size = 0
501
+
502
+ for i, sentence in enumerate(sentences):
503
+ sentence_size = count_units(sentence, unit)
504
+ if current_size + sentence_size > max_chunk_size and current_chunk:
505
+ chunks.append(' '.join(current_chunk))
506
+ # Use last 3 sentences for overlap
507
+ current_chunk = current_chunk[-3:]
508
+ current_size = count_units(' '.join(current_chunk), unit)
509
+
510
+ current_chunk.append(sentence)
511
+ current_size += sentence_size
512
+
513
+ if i + 1 < len(sentences):
514
+ current_vector = sentence_vectors[i]
515
+ next_vector = sentence_vectors[i + 1]
516
+ similarity = cosine_similarity(current_vector, next_vector)[0][0]
517
+ if similarity < 0.5 and current_size >= max_chunk_size // 2:
518
+ chunks.append(' '.join(current_chunk))
519
+ current_chunk = current_chunk[-3:]
520
+ current_size = count_units(' '.join(current_chunk), unit)
521
+
522
+ if current_chunk:
523
+ chunks.append(' '.join(current_chunk))
524
+
525
+ return chunks
526
+
527
+
528
+ def semantic_chunk_long_file(file_path: str, max_chunk_size: int = 1000, overlap: int = 100, unit: str = 'words') -> Optional[List[str]]:
529
+ logging.debug("semantic_chunk_long_file...")
530
+ try:
531
+ with open(file_path, 'r', encoding='utf-8') as file:
532
+ content = file.read()
533
+
534
+ chunks = semantic_chunking(content, max_chunk_size, unit)
535
+ return chunks
536
+ except Exception as e:
537
+ logging.error(f"Error chunking text file: {str(e)}")
538
+ return None
539
+
540
+ #
541
+ #
542
+ #######################################################################################################################
543
+
544
+
545
+ #######################################################################################################################
546
+ #
547
+ # Embedding Chunking
548
+
549
+ def chunk_for_embedding(text: str, file_name: str, custom_chunk_options: Dict[str, Any] = None) -> List[Dict[str, Any]]:
550
+ options = chunk_options.copy()
551
+ if custom_chunk_options:
552
+ options.update(custom_chunk_options)
553
+
554
+ logging.info(f"Chunking options: {options}")
555
+ chunks = improved_chunking_process(text, options)
556
+ total_chunks = len(chunks)
557
+ logging.info(f"Total chunks created: {total_chunks}")
558
+
559
+ chunked_text_with_headers = []
560
+ for i, chunk in enumerate(chunks, 1):
561
+ chunk_text = chunk['text']
562
+ chunk_position = determine_chunk_position(chunk['metadata']['relative_position'])
563
+ chunk_header = f"""
564
+ Original Document: {file_name}
565
+ Chunk: {i} of {total_chunks}
566
+ Position: {chunk_position}
567
+
568
+ --- Chunk Content ---
569
+ """
570
+
571
+ full_chunk_text = chunk_header + chunk_text
572
+ chunk['text'] = full_chunk_text
573
+ chunk['metadata']['file_name'] = file_name
574
+ chunked_text_with_headers.append(chunk)
575
+
576
+ return chunked_text_with_headers
577
+
578
+ #
579
+ # End of Embedding Chunking
580
+ #######################################################################################################################
581
+
582
+
583
+ #######################################################################################################################
584
+ #
585
+ # JSON Chunking
586
+
587
+ # FIXME
588
+ def chunk_text_by_json(text: str, max_size: int = 1000, overlap: int = 0) -> List[Dict[str, Any]]:
589
+ """
590
+ Chunk JSON-formatted text into smaller JSON chunks while preserving structure.
591
+
592
+ Parameters:
593
+ - text (str): The JSON-formatted text to be chunked.
594
+ - max_size (int): Maximum number of items or characters per chunk.
595
+ - overlap (int): Number of items or characters to overlap between chunks.
596
+
597
+ Returns:
598
+ - List[Dict[str, Any]]: A list of chunks with their metadata.
599
+ """
600
+ logging.debug("chunk_text_by_json started...")
601
+ try:
602
+ json_data = json.loads(text)
603
+ except json.JSONDecodeError as e:
604
+ logging.error(f"Invalid JSON data: {e}")
605
+ raise ValueError(f"Invalid JSON data: {e}")
606
+
607
+ # Determine if JSON data is a list or a dict
608
+ if isinstance(json_data, list):
609
+ return chunk_json_list(json_data, max_size, overlap)
610
+ elif isinstance(json_data, dict):
611
+ return chunk_json_dict(json_data, max_size, overlap)
612
+ else:
613
+ logging.error("Unsupported JSON structure. Only JSON objects and arrays are supported.")
614
+ raise ValueError("Unsupported JSON structure. Only JSON objects and arrays are supported.")
615
+
616
+
617
+ def chunk_json_list(json_list: List[Any], max_size: int, overlap: int) -> List[Dict[str, Any]]:
618
+ """
619
+ Chunk a JSON array into smaller chunks.
620
+
621
+ Parameters:
622
+ - json_list (List[Any]): The JSON array to be chunked.
623
+ - max_size (int): Maximum number of items per chunk.
624
+ - overlap (int): Number of items to overlap between chunks.
625
+
626
+ Returns:
627
+ - List[Dict[str, Any]]: A list of JSON chunks with metadata.
628
+ """
629
+ logging.debug("chunk_json_list started...")
630
+ chunks = []
631
+ total_items = len(json_list)
632
+ step = max_size - overlap
633
+ if step <= 0:
634
+ raise ValueError("max_size must be greater than overlap.")
635
+
636
+ for i in range(0, total_items, step):
637
+ chunk = json_list[i:i + max_size]
638
+ metadata = {
639
+ 'chunk_index': i // step + 1,
640
+ 'total_chunks': (total_items + step - 1) // step,
641
+ 'chunk_method': 'json_list',
642
+ 'max_size': max_size,
643
+ 'overlap': overlap,
644
+ 'relative_position': i / total_items
645
+ }
646
+ chunks.append({
647
+ 'json': chunk,
648
+ 'metadata': metadata
649
+ })
650
+
651
+ logging.debug(f"chunk_json_list created {len(chunks)} chunks.")
652
+ return chunks
653
+
654
+
655
+
656
+ def chunk_json_dict(json_dict: Dict[str, Any], max_size: int, overlap: int) -> List[Dict[str, Any]]:
657
+ """
658
+ Chunk a JSON object into smaller chunks based on its 'data' key while preserving other keys like 'metadata'.
659
+
660
+ Parameters:
661
+ - json_dict (Dict[str, Any]): The JSON object to be chunked.
662
+ - max_size (int): Maximum number of key-value pairs per chunk in the 'data' section.
663
+ - overlap (int): Number of key-value pairs to overlap between chunks.
664
+
665
+ Returns:
666
+ - List[Dict[str, Any]]: A list of JSON chunks with metadata.
667
+ """
668
+ logging.debug("chunk_json_dict started...")
669
+
670
+ # Preserve non-chunked sections
671
+ preserved_keys = ['metadata']
672
+ preserved_data = {key: value for key, value in json_dict.items() if key in preserved_keys}
673
+
674
+ # Identify the chunkable section
675
+ chunkable_key = 'data'
676
+ if chunkable_key not in json_dict or not isinstance(json_dict[chunkable_key], dict):
677
+ logging.error("No chunkable 'data' section found in JSON dictionary.")
678
+ raise ValueError("No chunkable 'data' section found in JSON dictionary.")
679
+
680
+ chunkable_data = json_dict[chunkable_key]
681
+ data_keys = list(chunkable_data.keys())
682
+ total_keys = len(data_keys)
683
+ chunks = []
684
+ step = max_size - overlap
685
+ if step <= 0:
686
+ raise ValueError("max_size must be greater than overlap.")
687
+
688
+ # Adjust the loop to prevent creating an extra chunk
689
+ for i in range(0, total_keys, step):
690
+ chunk_keys = data_keys[i:i + max_size]
691
+
692
+ # Handle overlap
693
+ if i != 0 and overlap > 0:
694
+ overlap_keys = data_keys[i - overlap:i]
695
+ chunk_keys = overlap_keys + chunk_keys
696
+
697
+ # Remove duplicate keys caused by overlap
698
+ unique_chunk_keys = []
699
+ seen_keys = set()
700
+ for key in chunk_keys:
701
+ if key not in seen_keys:
702
+ unique_chunk_keys.append(key)
703
+ seen_keys.add(key)
704
+
705
+ chunk_data = {key: chunkable_data[key] for key in unique_chunk_keys}
706
+
707
+ metadata = {
708
+ 'chunk_index': (i // step) + 1,
709
+ 'total_chunks': (total_keys + step - 1) // step,
710
+ 'chunk_method': 'json_dict',
711
+ 'max_size': max_size,
712
+ 'overlap': overlap,
713
+ 'language': 'english', # Assuming English; modify as needed
714
+ 'relative_position': (i // step + 1) / ((total_keys + step - 1) // step)
715
+ }
716
+
717
+ # Merge preserved data into metadata
718
+ metadata.update(preserved_data.get('metadata', {}))
719
+
720
+ # Create the chunk with preserved data
721
+ chunk = {
722
+ 'metadata': preserved_data,
723
+ 'data': chunk_data
724
+ }
725
+
726
+ chunks.append({
727
+ 'json': chunk,
728
+ 'metadata': metadata
729
+ })
730
+
731
+ logging.debug(f"chunk_json_dict created {len(chunks)} chunks.")
732
+ return chunks
733
+
734
+
735
+ #
736
+ # End of JSON Chunking
737
+ #######################################################################################################################
738
+
739
+ #######################################################################################################################
740
+ #
741
+ # OpenAI Rolling Summarization
742
+ #
743
+
744
+ client = OpenAI(api_key=openai_api_key)
745
+ def get_chat_completion(messages, model='gpt-4-turbo'):
746
+ response = client.chat.completions.create(
747
+ model=model,
748
+ messages=messages,
749
+ temperature=0,
750
+ )
751
+ return response.choices[0].message.content
752
+
753
+
754
+ # This function combines text chunks into larger blocks without exceeding a specified token count.
755
+ # It returns the combined chunks, their original indices, and the number of dropped chunks due to overflow.
756
+ def combine_chunks_with_no_minimum(
757
+ chunks: List[str],
758
+ max_tokens: int,
759
+ chunk_delimiter: str = "\n\n",
760
+ header: Optional[str] = None,
761
+ add_ellipsis_for_overflow: bool = False,
762
+ ) -> Tuple[List[str], List[List[int]], int]:
763
+ dropped_chunk_count = 0
764
+ output = [] # list to hold the final combined chunks
765
+ output_indices = [] # list to hold the indices of the final combined chunks
766
+ candidate = [header] if header else [] # list to hold the current combined chunk candidate
767
+ candidate_indices = []
768
+ for chunk_i, chunk in enumerate(chunks):
769
+ chunk_with_header = [chunk] if not header else [header, chunk]
770
+ combined_text = chunk_delimiter.join(candidate + chunk_with_header)
771
+ token_count = len(tokenizer.encode(combined_text))
772
+ if token_count > max_tokens:
773
+ if add_ellipsis_for_overflow and len(candidate) > 0:
774
+ ellipsis_text = chunk_delimiter.join(candidate + ["..."])
775
+ if len(tokenizer.encode(ellipsis_text)) <= max_tokens:
776
+ candidate = candidate + ["..."]
777
+ dropped_chunk_count += 1
778
+ if len(candidate) > 0:
779
+ output.append(chunk_delimiter.join(candidate))
780
+ output_indices.append(candidate_indices)
781
+ candidate = chunk_with_header
782
+ candidate_indices = [chunk_i]
783
+ else:
784
+ logging.warning(f"Single chunk at index {chunk_i} exceeds max_tokens and will be dropped.")
785
+ dropped_chunk_count += 1
786
+ else:
787
+ candidate.extend(chunk_with_header)
788
+ candidate_indices.append(chunk_i)
789
+
790
+ if candidate:
791
+ output.append(chunk_delimiter.join(candidate))
792
+ output_indices.append(candidate_indices)
793
+ return output, output_indices, dropped_chunk_count
794
+
795
+
796
+ def rolling_summarize(text: str,
797
+ detail: float = 0,
798
+ model: str = 'gpt-4o',
799
+ additional_instructions: Optional[str] = None,
800
+ minimum_chunk_size: Optional[int] = 500,
801
+ chunk_delimiter: str = ".",
802
+ summarize_recursively: bool = False,
803
+ verbose: bool = False) -> str:
804
+ """
805
+ Summarizes a given text by splitting it into chunks, each of which is summarized individually.
806
+ The level of detail in the summary can be adjusted, and the process can optionally be made recursive.
807
+
808
+ Parameters:
809
+ - text (str): The text to be summarized.
810
+ - detail (float, optional): A value between 0 and 1 indicating the desired level of detail in the summary.
811
+ - additional_instructions (Optional[str], optional): Additional instructions for the model.
812
+ - minimum_chunk_size (Optional[int], optional): The minimum size for text chunks.
813
+ - chunk_delimiter (str, optional): The delimiter used to split the text into chunks.
814
+ - summarize_recursively (bool, optional): If True, summaries are generated recursively.
815
+ - verbose (bool, optional): If True, prints detailed information about the chunking process.
816
+
817
+ Returns:
818
+ - str: The final compiled summary of the text.
819
+
820
+ The function first determines the number of chunks by interpolating between a minimum and a maximum chunk count
821
+ based on the `detail` parameter. It then splits the text into chunks and summarizes each chunk. If
822
+ `summarize_recursively` is True, each summary is based on the previous summaries, adding more context to the
823
+ summarization process. The function returns a compiled summary of all chunks.
824
+ """
825
+
826
+ # Check detail is set correctly
827
+ assert 0 <= detail <= 1, "Detail must be between 0 and 1."
828
+
829
+ # Interpolate the number of chunks based on the detail parameter
830
+ text_length = len(tokenizer.encode(text))
831
+ max_chunks = text_length // minimum_chunk_size if minimum_chunk_size else 10
832
+ min_chunks = 1
833
+ num_chunks = int(min_chunks + detail * (max_chunks - min_chunks))
834
+
835
+ # Adjust chunk_size based on interpolated number of chunks
836
+ chunk_size = max(minimum_chunk_size, text_length // num_chunks) if num_chunks else text_length
837
+ text_chunks = chunk_on_delimiter(text, chunk_size, chunk_delimiter)
838
+ if verbose:
839
+ print(f"Splitting the text into {len(text_chunks)} chunks to be summarized.")
840
+ print(f"Chunk lengths are {[len(tokenizer.encode(x)) for x in text_chunks]} tokens.")
841
+
842
+ # Set system message
843
+ system_message_content = "Rewrite this text in summarized form."
844
+ if additional_instructions:
845
+ system_message_content += f"\n\n{additional_instructions}"
846
+
847
+ accumulated_summaries = []
848
+ for i, chunk in enumerate(tqdm(text_chunks, desc="Summarizing chunks")):
849
+ if summarize_recursively and accumulated_summaries:
850
+ # Combine previous summary with current chunk for recursive summarization
851
+ combined_text = accumulated_summaries[-1] + "\n\n" + chunk
852
+ user_message_content = f"Previous summary and new content to summarize:\n\n{combined_text}"
853
+ else:
854
+ user_message_content = chunk
855
+
856
+ messages = [
857
+ {"role": "system", "content": system_message_content},
858
+ {"role": "user", "content": user_message_content}
859
+ ]
860
+
861
+ response = get_chat_completion(messages, model=model)
862
+ accumulated_summaries.append(response)
863
+
864
+ final_summary = '\n\n'.join(accumulated_summaries)
865
+ return final_summary
866
+
867
+ #
868
+ #
869
+ #######################################################################################################################
870
+ #
871
+ # Ebook Chapter Chunking
872
+
873
+
874
+ def chunk_ebook_by_chapters(text: str, chunk_options: Dict[str, Any]) -> List[Dict[str, Any]]:
875
+ logging.debug("chunk_ebook_by_chapters")
876
+ max_chunk_size = int(chunk_options.get('max_size', 300))
877
+ overlap = int(chunk_options.get('overlap', 0))
878
+ custom_pattern = chunk_options.get('custom_chapter_pattern', None)
879
+
880
+ # List of chapter heading patterns to try, in order
881
+ chapter_patterns = [
882
+ custom_pattern,
883
+ r'^#{1,2}\s+', # Markdown style: '# ' or '## '
884
+ r'^Chapter\s+\d+', # 'Chapter ' followed by numbers
885
+ r'^\d+\.\s+', # Numbered chapters: '1. ', '2. ', etc.
886
+ r'^[A-Z\s]+$' # All caps headings
887
+ ]
888
+
889
+ chapter_positions = []
890
+ used_pattern = None
891
+
892
+ for pattern in chapter_patterns:
893
+ if pattern is None:
894
+ continue
895
+ chapter_regex = re.compile(pattern, re.MULTILINE | re.IGNORECASE)
896
+ chapter_positions = [match.start() for match in chapter_regex.finditer(text)]
897
+ if chapter_positions:
898
+ used_pattern = pattern
899
+ break
900
+
901
+ # If no chapters found, return the entire content as one chunk
902
+ if not chapter_positions:
903
+ metadata = get_chunk_metadata(
904
+ chunk=text,
905
+ full_text=text,
906
+ chunk_type="whole_document",
907
+ language=chunk_options.get('language', 'english')
908
+ )
909
+ return [{'text': text, 'metadata': metadata}]
910
+
911
+ # Split content into chapters
912
+ chunks = []
913
+ for i in range(len(chapter_positions)):
914
+ start = chapter_positions[i]
915
+ end = chapter_positions[i + 1] if i + 1 < len(chapter_positions) else None
916
+ chapter = text[start:end]
917
+
918
+ # Apply overlap if specified
919
+ if overlap > 0 and i > 0:
920
+ overlap_start = max(0, chapter_positions[i] - overlap)
921
+ chapter = text[overlap_start:end]
922
+
923
+ chunks.append(chapter)
924
+
925
+ # Post-process chunks
926
+ processed_chunks = post_process_chunks(chunks)
927
+
928
+ # Add metadata to chunks
929
+ chunks_with_metadata = []
930
+ for i, chunk in enumerate(processed_chunks):
931
+ metadata = get_chunk_metadata(
932
+ chunk=chunk,
933
+ full_text=text,
934
+ chunk_type="chapter",
935
+ chapter_number=i + 1,
936
+ chapter_pattern=used_pattern,
937
+ language=chunk_options.get('language', 'english')
938
+ )
939
+ chunks_with_metadata.append({'text': chunk, 'metadata': metadata})
940
+
941
+ return chunks_with_metadata
942
+
943
+ #
944
+ # End of ebook chapter chunking
945
+ #######################################################################################################################
946
+
947
+ #######################################################################################################################
948
+ #
949
+ # Functions for adapative chunking:
950
+
951
+ # FIXME - punkt
952
+
953
+ def adaptive_chunk_size(text: str, base_size: int = 1000, min_size: int = 500, max_size: int = 2000) -> int:
954
+ # Tokenize the text into sentences
955
+ sentences = sent_tokenize(text)
956
+
957
+ if not sentences:
958
+ return base_size
959
+
960
+ # Calculate average sentence length
961
+ avg_sentence_length = sum(len(s.split()) for s in sentences) / len(sentences)
962
+
963
+ # Adjust chunk size based on average sentence length
964
+ if avg_sentence_length < 10:
965
+ size_factor = 1.2 # Increase chunk size for short sentences
966
+ elif avg_sentence_length > 20:
967
+ size_factor = 0.8 # Decrease chunk size for long sentences
968
+ else:
969
+ size_factor = 1.0
970
+
971
+ # Calculate adaptive chunk size
972
+ adaptive_size = int(base_size * size_factor)
973
+
974
+ # Ensure chunk size is within bounds
975
+ return max(min_size, min(adaptive_size, max_size))
976
+
977
+
978
+ def adaptive_chunk_size_non_punkt(text: str, base_size: int, min_size: int = 100, max_size: int = 2000) -> int:
979
+ # Adaptive logic: adjust chunk size based on text complexity
980
+ words = text.split()
981
+ if not words:
982
+ return base_size # Return base_size if text is empty
983
+
984
+ avg_word_length = sum(len(word) for word in words) / len(words)
985
+
986
+ if avg_word_length > 6: # Threshold for "complex" text
987
+ adjusted_size = int(base_size * 0.8) # Reduce chunk size for complex text
988
+ elif avg_word_length < 4: # Threshold for "simple" text
989
+ adjusted_size = int(base_size * 1.2) # Increase chunk size for simple text
990
+ else:
991
+ adjusted_size = base_size
992
+
993
+ # Ensure the chunk size is within the specified range
994
+ return max(min_size, min(adjusted_size, max_size))
995
+
996
+
997
+ def adaptive_chunking(text: str, base_size: int = 1000, min_size: int = 500, max_size: int = 2000) -> List[str]:
998
+ logging.debug("adaptive_chunking...")
999
+ chunk_size = adaptive_chunk_size(text, base_size, min_size, max_size)
1000
+ words = text.split()
1001
+ chunks = []
1002
+ current_chunk = []
1003
+ current_length = 0
1004
+
1005
+ for word in words:
1006
+ if current_length + len(word) > chunk_size and current_chunk:
1007
+ chunks.append(' '.join(current_chunk))
1008
+ current_chunk = []
1009
+ current_length = 0
1010
+ current_chunk.append(word)
1011
+ current_length += len(word) + 1 # +1 for space
1012
+
1013
+ if current_chunk:
1014
+ chunks.append(' '.join(current_chunk))
1015
+
1016
+ return chunks
1017
+
1018
+ # FIXME - usage example
1019
+ # chunk_options = {
1020
+ # 'method': 'words', # or any other method
1021
+ # 'base_size': 1000,
1022
+ # 'min_size': 100,
1023
+ # 'max_size': 2000,
1024
+ # 'adaptive': True,
1025
+ # 'language': 'en'
1026
+ # }
1027
+ #chunks = improved_chunking_process(your_text, chunk_options)
1028
+
1029
+
1030
+ # Example of chunking a document with metadata
1031
+ # document_metadata = {
1032
+ # 'title': 'Example Document',
1033
+ # 'author': 'John Doe',
1034
+ # 'creation_date': '2023-06-14',
1035
+ # 'source': 'https://example.com/document',
1036
+ # 'document_type': 'article'
1037
+ # }
1038
+ #
1039
+ # chunk_options = {
1040
+ # 'method': 'sentences',
1041
+ # 'base_size': 1000,
1042
+ # 'adaptive': True,
1043
+ # 'language': 'en'
1044
+ # }
1045
+ #
1046
+ # processed_document = process_document_with_metadata(your_text, chunk_options, document_metadata)
1047
+
1048
+
1049
+ #
1050
+ # End of Chunking Library
1051
+ #######################################################################################################################
App_Function_Libraries/DB/Character_Chat_DB.py ADDED
@@ -0,0 +1,701 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # character_chat_db.py
2
+ # Database functions for managing character cards and chat histories.
3
+ # #
4
+ # Imports
5
+ import configparser
6
+ import sqlite3
7
+ import json
8
+ import os
9
+ import sys
10
+ from typing import List, Dict, Optional, Tuple, Any, Union
11
+
12
+ from App_Function_Libraries.Utils.Utils import get_database_dir, get_project_relative_path, get_database_path
13
+ from Tests.Chat_APIs.Chat_APIs_Integration_test import logging
14
+
15
+ #
16
+ #######################################################################################################################
17
+ #
18
+ #
19
+
20
+ def ensure_database_directory():
21
+ os.makedirs(get_database_dir(), exist_ok=True)
22
+
23
+ ensure_database_directory()
24
+
25
+
26
+ # Construct the path to the config file
27
+ config_path = get_project_relative_path('Config_Files/config.txt')
28
+
29
+ # Read the config file
30
+ config = configparser.ConfigParser()
31
+ config.read(config_path)
32
+
33
+ # Get the chat db path from the config, or use the default if not specified
34
+ chat_DB_PATH = config.get('Database', 'chatDB_path', fallback=get_database_path('chatDB.db'))
35
+ print(f"Chat Database path: {chat_DB_PATH}")
36
+
37
+ ########################################################################################################
38
+ #
39
+ # Functions
40
+
41
+ # FIXME - Setup properly and test/add documentation for its existence...
42
+ def initialize_database():
43
+ """Initialize the SQLite database with required tables and FTS5 virtual tables."""
44
+ conn = None
45
+ try:
46
+ conn = sqlite3.connect(chat_DB_PATH)
47
+ cursor = conn.cursor()
48
+
49
+ # Enable foreign key constraints
50
+ cursor.execute("PRAGMA foreign_keys = ON;")
51
+
52
+ # Create CharacterCards table with V2 fields
53
+ cursor.execute("""
54
+ CREATE TABLE IF NOT EXISTS CharacterCards (
55
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
56
+ name TEXT UNIQUE NOT NULL,
57
+ description TEXT,
58
+ personality TEXT,
59
+ scenario TEXT,
60
+ image BLOB,
61
+ post_history_instructions TEXT,
62
+ first_mes TEXT,
63
+ mes_example TEXT,
64
+ creator_notes TEXT,
65
+ system_prompt TEXT,
66
+ alternate_greetings TEXT,
67
+ tags TEXT,
68
+ creator TEXT,
69
+ character_version TEXT,
70
+ extensions TEXT,
71
+ created_at DATETIME DEFAULT CURRENT_TIMESTAMP
72
+ );
73
+ """)
74
+
75
+ # Create CharacterChats table
76
+ cursor.execute("""
77
+ CREATE TABLE IF NOT EXISTS CharacterChats (
78
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
79
+ character_id INTEGER NOT NULL,
80
+ conversation_name TEXT,
81
+ chat_history TEXT,
82
+ is_snapshot BOOLEAN DEFAULT FALSE,
83
+ created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
84
+ FOREIGN KEY (character_id) REFERENCES CharacterCards(id) ON DELETE CASCADE
85
+ );
86
+ """)
87
+
88
+ # Create FTS5 virtual table for CharacterChats
89
+ cursor.execute("""
90
+ CREATE VIRTUAL TABLE IF NOT EXISTS CharacterChats_fts USING fts5(
91
+ conversation_name,
92
+ chat_history,
93
+ content='CharacterChats',
94
+ content_rowid='id'
95
+ );
96
+ """)
97
+
98
+ # Create triggers to keep FTS5 table in sync with CharacterChats
99
+ cursor.executescript("""
100
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_ai AFTER INSERT ON CharacterChats BEGIN
101
+ INSERT INTO CharacterChats_fts(rowid, conversation_name, chat_history)
102
+ VALUES (new.id, new.conversation_name, new.chat_history);
103
+ END;
104
+
105
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_ad AFTER DELETE ON CharacterChats BEGIN
106
+ DELETE FROM CharacterChats_fts WHERE rowid = old.id;
107
+ END;
108
+
109
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_au AFTER UPDATE ON CharacterChats BEGIN
110
+ UPDATE CharacterChats_fts SET conversation_name = new.conversation_name, chat_history = new.chat_history
111
+ WHERE rowid = new.id;
112
+ END;
113
+ """)
114
+
115
+ # Create ChatKeywords table
116
+ cursor.execute("""
117
+ CREATE TABLE IF NOT EXISTS ChatKeywords (
118
+ chat_id INTEGER NOT NULL,
119
+ keyword TEXT NOT NULL,
120
+ FOREIGN KEY (chat_id) REFERENCES CharacterChats(id) ON DELETE CASCADE
121
+ );
122
+ """)
123
+
124
+ # Create indexes for faster searches
125
+ cursor.execute("""
126
+ CREATE INDEX IF NOT EXISTS idx_chatkeywords_keyword ON ChatKeywords(keyword);
127
+ """)
128
+ cursor.execute("""
129
+ CREATE INDEX IF NOT EXISTS idx_chatkeywords_chat_id ON ChatKeywords(chat_id);
130
+ """)
131
+
132
+ conn.commit()
133
+ logging.info("Database initialized successfully.")
134
+ except sqlite3.Error as e:
135
+ logging.error(f"SQLite error occurred during database initialization: {e}")
136
+ if conn:
137
+ conn.rollback()
138
+ raise
139
+ except Exception as e:
140
+ logging.error(f"Unexpected error occurred during database initialization: {e}")
141
+ if conn:
142
+ conn.rollback()
143
+ raise
144
+ finally:
145
+ if conn:
146
+ conn.close()
147
+
148
+ # Call initialize_database() at the start of your application
149
+ def setup_chat_database():
150
+ try:
151
+ initialize_database()
152
+ except Exception as e:
153
+ logging.critical(f"Failed to initialize database: {e}")
154
+ sys.exit(1)
155
+
156
+ setup_chat_database()
157
+
158
+ ########################################################################################################
159
+ #
160
+ # Character Card handling
161
+
162
+ def parse_character_card(card_data: Dict[str, Any]) -> Dict[str, Any]:
163
+ """Parse and validate a character card according to V2 specification."""
164
+ v2_data = {
165
+ 'name': card_data.get('name', ''),
166
+ 'description': card_data.get('description', ''),
167
+ 'personality': card_data.get('personality', ''),
168
+ 'scenario': card_data.get('scenario', ''),
169
+ 'first_mes': card_data.get('first_mes', ''),
170
+ 'mes_example': card_data.get('mes_example', ''),
171
+ 'creator_notes': card_data.get('creator_notes', ''),
172
+ 'system_prompt': card_data.get('system_prompt', ''),
173
+ 'post_history_instructions': card_data.get('post_history_instructions', ''),
174
+ 'alternate_greetings': json.dumps(card_data.get('alternate_greetings', [])),
175
+ 'tags': json.dumps(card_data.get('tags', [])),
176
+ 'creator': card_data.get('creator', ''),
177
+ 'character_version': card_data.get('character_version', ''),
178
+ 'extensions': json.dumps(card_data.get('extensions', {}))
179
+ }
180
+
181
+ # Handle 'image' separately as it might be binary data
182
+ if 'image' in card_data:
183
+ v2_data['image'] = card_data['image']
184
+
185
+ return v2_data
186
+
187
+
188
+ def add_character_card(card_data: Dict[str, Any]) -> Optional[int]:
189
+ """Add or update a character card in the database."""
190
+ conn = sqlite3.connect(chat_DB_PATH)
191
+ cursor = conn.cursor()
192
+ try:
193
+ parsed_card = parse_character_card(card_data)
194
+
195
+ # Check if character already exists
196
+ cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (parsed_card['name'],))
197
+ row = cursor.fetchone()
198
+
199
+ if row:
200
+ # Update existing character
201
+ character_id = row[0]
202
+ update_query = """
203
+ UPDATE CharacterCards
204
+ SET description = ?, personality = ?, scenario = ?, image = ?,
205
+ post_history_instructions = ?, first_mes = ?, mes_example = ?,
206
+ creator_notes = ?, system_prompt = ?, alternate_greetings = ?,
207
+ tags = ?, creator = ?, character_version = ?, extensions = ?
208
+ WHERE id = ?
209
+ """
210
+ cursor.execute(update_query, (
211
+ parsed_card['description'], parsed_card['personality'], parsed_card['scenario'],
212
+ parsed_card['image'], parsed_card['post_history_instructions'], parsed_card['first_mes'],
213
+ parsed_card['mes_example'], parsed_card['creator_notes'], parsed_card['system_prompt'],
214
+ parsed_card['alternate_greetings'], parsed_card['tags'], parsed_card['creator'],
215
+ parsed_card['character_version'], parsed_card['extensions'], character_id
216
+ ))
217
+ else:
218
+ # Insert new character
219
+ insert_query = """
220
+ INSERT INTO CharacterCards (name, description, personality, scenario, image,
221
+ post_history_instructions, first_mes, mes_example, creator_notes, system_prompt,
222
+ alternate_greetings, tags, creator, character_version, extensions)
223
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
224
+ """
225
+ cursor.execute(insert_query, (
226
+ parsed_card['name'], parsed_card['description'], parsed_card['personality'],
227
+ parsed_card['scenario'], parsed_card['image'], parsed_card['post_history_instructions'],
228
+ parsed_card['first_mes'], parsed_card['mes_example'], parsed_card['creator_notes'],
229
+ parsed_card['system_prompt'], parsed_card['alternate_greetings'], parsed_card['tags'],
230
+ parsed_card['creator'], parsed_card['character_version'], parsed_card['extensions']
231
+ ))
232
+ character_id = cursor.lastrowid
233
+
234
+ conn.commit()
235
+ return character_id
236
+ except sqlite3.IntegrityError as e:
237
+ logging.error(f"Error adding character card: {e}")
238
+ return None
239
+ except Exception as e:
240
+ logging.error(f"Unexpected error adding character card: {e}")
241
+ return None
242
+ finally:
243
+ conn.close()
244
+
245
+ # def add_character_card(card_data: Dict) -> Optional[int]:
246
+ # """Add or update a character card in the database.
247
+ #
248
+ # Returns the ID of the inserted character or None if failed.
249
+ # """
250
+ # conn = sqlite3.connect(chat_DB_PATH)
251
+ # cursor = conn.cursor()
252
+ # try:
253
+ # # Ensure all required fields are present
254
+ # required_fields = ['name', 'description', 'personality', 'scenario', 'image', 'post_history_instructions', 'first_message']
255
+ # for field in required_fields:
256
+ # if field not in card_data:
257
+ # card_data[field] = '' # Assign empty string if field is missing
258
+ #
259
+ # # Check if character already exists
260
+ # cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (card_data['name'],))
261
+ # row = cursor.fetchone()
262
+ #
263
+ # if row:
264
+ # # Update existing character
265
+ # character_id = row[0]
266
+ # cursor.execute("""
267
+ # UPDATE CharacterCards
268
+ # SET description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ?
269
+ # WHERE id = ?
270
+ # """, (
271
+ # card_data['description'],
272
+ # card_data['personality'],
273
+ # card_data['scenario'],
274
+ # card_data['image'],
275
+ # card_data['post_history_instructions'],
276
+ # card_data['first_message'],
277
+ # character_id
278
+ # ))
279
+ # else:
280
+ # # Insert new character
281
+ # cursor.execute("""
282
+ # INSERT INTO CharacterCards (name, description, personality, scenario, image, post_history_instructions, first_message)
283
+ # VALUES (?, ?, ?, ?, ?, ?, ?)
284
+ # """, (
285
+ # card_data['name'],
286
+ # card_data['description'],
287
+ # card_data['personality'],
288
+ # card_data['scenario'],
289
+ # card_data['image'],
290
+ # card_data['post_history_instructions'],
291
+ # card_data['first_message']
292
+ # ))
293
+ # character_id = cursor.lastrowid
294
+ #
295
+ # conn.commit()
296
+ # return cursor.lastrowid
297
+ # except sqlite3.IntegrityError as e:
298
+ # logging.error(f"Error adding character card: {e}")
299
+ # return None
300
+ # except Exception as e:
301
+ # logging.error(f"Unexpected error adding character card: {e}")
302
+ # return None
303
+ # finally:
304
+ # conn.close()
305
+
306
+
307
+ def get_character_cards() -> List[Dict]:
308
+ """Retrieve all character cards from the database."""
309
+ logging.debug(f"Fetching characters from DB: {chat_DB_PATH}")
310
+ conn = sqlite3.connect(chat_DB_PATH)
311
+ cursor = conn.cursor()
312
+ cursor.execute("SELECT * FROM CharacterCards")
313
+ rows = cursor.fetchall()
314
+ columns = [description[0] for description in cursor.description]
315
+ conn.close()
316
+ characters = [dict(zip(columns, row)) for row in rows]
317
+ #logging.debug(f"Characters fetched from DB: {characters}")
318
+ return characters
319
+
320
+
321
+ def get_character_card_by_id(character_id: Union[int, Dict[str, Any]]) -> Optional[Dict[str, Any]]:
322
+ """
323
+ Retrieve a single character card by its ID.
324
+
325
+ Args:
326
+ character_id: Can be either an integer ID or a dictionary containing character data.
327
+
328
+ Returns:
329
+ A dictionary containing the character card data, or None if not found.
330
+ """
331
+ conn = sqlite3.connect(chat_DB_PATH)
332
+ cursor = conn.cursor()
333
+ try:
334
+ if isinstance(character_id, dict):
335
+ # If a dictionary is passed, assume it's already a character card
336
+ return character_id
337
+ elif isinstance(character_id, int):
338
+ # If an integer is passed, fetch the character from the database
339
+ cursor.execute("SELECT * FROM CharacterCards WHERE id = ?", (character_id,))
340
+ row = cursor.fetchone()
341
+ if row:
342
+ columns = [description[0] for description in cursor.description]
343
+ return dict(zip(columns, row))
344
+ else:
345
+ logging.warning(f"Invalid type for character_id: {type(character_id)}")
346
+ return None
347
+ except Exception as e:
348
+ logging.error(f"Error in get_character_card_by_id: {e}")
349
+ return None
350
+ finally:
351
+ conn.close()
352
+
353
+
354
+ def update_character_card(character_id: int, card_data: Dict) -> bool:
355
+ """Update an existing character card."""
356
+ conn = sqlite3.connect(chat_DB_PATH)
357
+ cursor = conn.cursor()
358
+ try:
359
+ cursor.execute("""
360
+ UPDATE CharacterCards
361
+ SET name = ?, description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ?
362
+ WHERE id = ?
363
+ """, (
364
+ card_data.get('name'),
365
+ card_data.get('description'),
366
+ card_data.get('personality'),
367
+ card_data.get('scenario'),
368
+ card_data.get('image'),
369
+ card_data.get('post_history_instructions', ''),
370
+ card_data.get('first_message', "Hello! I'm ready to chat."),
371
+ character_id
372
+ ))
373
+ conn.commit()
374
+ return cursor.rowcount > 0
375
+ except sqlite3.IntegrityError as e:
376
+ logging.error(f"Error updating character card: {e}")
377
+ return False
378
+ finally:
379
+ conn.close()
380
+
381
+
382
+ def delete_character_card(character_id: int) -> bool:
383
+ """Delete a character card and its associated chats."""
384
+ conn = sqlite3.connect(chat_DB_PATH)
385
+ cursor = conn.cursor()
386
+ try:
387
+ # Delete associated chats first due to foreign key constraint
388
+ cursor.execute("DELETE FROM CharacterChats WHERE character_id = ?", (character_id,))
389
+ cursor.execute("DELETE FROM CharacterCards WHERE id = ?", (character_id,))
390
+ conn.commit()
391
+ return cursor.rowcount > 0
392
+ except sqlite3.Error as e:
393
+ logging.error(f"Error deleting character card: {e}")
394
+ return False
395
+ finally:
396
+ conn.close()
397
+
398
+
399
+ def add_character_chat(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]], keywords: Optional[List[str]] = None, is_snapshot: bool = False) -> Optional[int]:
400
+ """
401
+ Add a new chat history for a character, optionally associating keywords.
402
+
403
+ Args:
404
+ character_id (int): The ID of the character.
405
+ conversation_name (str): Name of the conversation.
406
+ chat_history (List[Tuple[str, str]]): List of (user, bot) message tuples.
407
+ keywords (Optional[List[str]]): List of keywords to associate with this chat.
408
+ is_snapshot (bool, optional): Whether this chat is a snapshot.
409
+
410
+ Returns:
411
+ Optional[int]: The ID of the inserted chat or None if failed.
412
+ """
413
+ conn = sqlite3.connect(chat_DB_PATH)
414
+ cursor = conn.cursor()
415
+ try:
416
+ chat_history_json = json.dumps(chat_history)
417
+ cursor.execute("""
418
+ INSERT INTO CharacterChats (character_id, conversation_name, chat_history, is_snapshot)
419
+ VALUES (?, ?, ?, ?)
420
+ """, (
421
+ character_id,
422
+ conversation_name,
423
+ chat_history_json,
424
+ is_snapshot
425
+ ))
426
+ chat_id = cursor.lastrowid
427
+
428
+ if keywords:
429
+ # Insert keywords into ChatKeywords table
430
+ keyword_records = [(chat_id, keyword.strip().lower()) for keyword in keywords]
431
+ cursor.executemany("""
432
+ INSERT INTO ChatKeywords (chat_id, keyword)
433
+ VALUES (?, ?)
434
+ """, keyword_records)
435
+
436
+ conn.commit()
437
+ return chat_id
438
+ except sqlite3.Error as e:
439
+ logging.error(f"Error adding character chat: {e}")
440
+ return None
441
+ finally:
442
+ conn.close()
443
+
444
+
445
+ def get_character_chats(character_id: Optional[int] = None) -> List[Dict]:
446
+ """Retrieve all chats, or chats for a specific character if character_id is provided."""
447
+ conn = sqlite3.connect(chat_DB_PATH)
448
+ cursor = conn.cursor()
449
+ if character_id is not None:
450
+ cursor.execute("SELECT * FROM CharacterChats WHERE character_id = ?", (character_id,))
451
+ else:
452
+ cursor.execute("SELECT * FROM CharacterChats")
453
+ rows = cursor.fetchall()
454
+ columns = [description[0] for description in cursor.description]
455
+ conn.close()
456
+ return [dict(zip(columns, row)) for row in rows]
457
+
458
+
459
+ def get_character_chat_by_id(chat_id: int) -> Optional[Dict]:
460
+ """Retrieve a single chat by its ID."""
461
+ conn = sqlite3.connect(chat_DB_PATH)
462
+ cursor = conn.cursor()
463
+ cursor.execute("SELECT * FROM CharacterChats WHERE id = ?", (chat_id,))
464
+ row = cursor.fetchone()
465
+ conn.close()
466
+ if row:
467
+ columns = [description[0] for description in cursor.description]
468
+ chat = dict(zip(columns, row))
469
+ chat['chat_history'] = json.loads(chat['chat_history'])
470
+ return chat
471
+ return None
472
+
473
+
474
+ def search_character_chats(query: str, character_id: Optional[int] = None) -> Tuple[List[Dict], str]:
475
+ """
476
+ Search for character chats using FTS5, optionally filtered by character_id.
477
+
478
+ Args:
479
+ query (str): The search query.
480
+ character_id (Optional[int]): The ID of the character to filter chats by.
481
+
482
+ Returns:
483
+ Tuple[List[Dict], str]: A list of matching chats and a status message.
484
+ """
485
+ if not query.strip():
486
+ return [], "Please enter a search query."
487
+
488
+ conn = sqlite3.connect(chat_DB_PATH)
489
+ cursor = conn.cursor()
490
+ try:
491
+ if character_id is not None:
492
+ # Search with character_id filter
493
+ cursor.execute("""
494
+ SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history
495
+ FROM CharacterChats_fts
496
+ JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id
497
+ WHERE CharacterChats_fts MATCH ? AND CharacterChats.character_id = ?
498
+ ORDER BY rank
499
+ """, (query, character_id))
500
+ else:
501
+ # Search without character_id filter
502
+ cursor.execute("""
503
+ SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history
504
+ FROM CharacterChats_fts
505
+ JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id
506
+ WHERE CharacterChats_fts MATCH ?
507
+ ORDER BY rank
508
+ """, (query,))
509
+
510
+ rows = cursor.fetchall()
511
+ columns = [description[0] for description in cursor.description]
512
+ results = [dict(zip(columns, row)) for row in rows]
513
+
514
+ if character_id is not None:
515
+ status_message = f"Found {len(results)} chat(s) matching '{query}' for the selected character."
516
+ else:
517
+ status_message = f"Found {len(results)} chat(s) matching '{query}' across all characters."
518
+
519
+ return results, status_message
520
+ except Exception as e:
521
+ logging.error(f"Error searching chats with FTS5: {e}")
522
+ return [], f"Error occurred during search: {e}"
523
+ finally:
524
+ conn.close()
525
+
526
+ def update_character_chat(chat_id: int, chat_history: List[Tuple[str, str]]) -> bool:
527
+ """Update an existing chat history."""
528
+ conn = sqlite3.connect(chat_DB_PATH)
529
+ cursor = conn.cursor()
530
+ try:
531
+ chat_history_json = json.dumps(chat_history)
532
+ cursor.execute("""
533
+ UPDATE CharacterChats
534
+ SET chat_history = ?
535
+ WHERE id = ?
536
+ """, (
537
+ chat_history_json,
538
+ chat_id
539
+ ))
540
+ conn.commit()
541
+ return cursor.rowcount > 0
542
+ except sqlite3.Error as e:
543
+ logging.error(f"Error updating character chat: {e}")
544
+ return False
545
+ finally:
546
+ conn.close()
547
+
548
+
549
+ def delete_character_chat(chat_id: int) -> bool:
550
+ """Delete a specific chat."""
551
+ conn = sqlite3.connect(chat_DB_PATH)
552
+ cursor = conn.cursor()
553
+ try:
554
+ cursor.execute("DELETE FROM CharacterChats WHERE id = ?", (chat_id,))
555
+ conn.commit()
556
+ return cursor.rowcount > 0
557
+ except sqlite3.Error as e:
558
+ logging.error(f"Error deleting character chat: {e}")
559
+ return False
560
+ finally:
561
+ conn.close()
562
+
563
+ def fetch_keywords_for_chats(keywords: List[str]) -> List[int]:
564
+ """
565
+ Fetch chat IDs associated with any of the specified keywords.
566
+
567
+ Args:
568
+ keywords (List[str]): List of keywords to search for.
569
+
570
+ Returns:
571
+ List[int]: List of chat IDs associated with the keywords.
572
+ """
573
+ if not keywords:
574
+ return []
575
+
576
+ conn = sqlite3.connect(chat_DB_PATH)
577
+ cursor = conn.cursor()
578
+ try:
579
+ # Construct the WHERE clause to search for each keyword
580
+ keyword_clauses = " OR ".join(["keyword = ?"] * len(keywords))
581
+ sql_query = f"SELECT DISTINCT chat_id FROM ChatKeywords WHERE {keyword_clauses}"
582
+ cursor.execute(sql_query, keywords)
583
+ rows = cursor.fetchall()
584
+ chat_ids = [row[0] for row in rows]
585
+ return chat_ids
586
+ except Exception as e:
587
+ logging.error(f"Error in fetch_keywords_for_chats: {e}")
588
+ return []
589
+ finally:
590
+ conn.close()
591
+
592
+ def save_chat_history_to_character_db(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]]) -> Optional[int]:
593
+ """Save chat history to the CharacterChats table.
594
+
595
+ Returns the ID of the inserted chat or None if failed.
596
+ """
597
+ return add_character_chat(character_id, conversation_name, chat_history)
598
+
599
+ def migrate_chat_to_media_db():
600
+ pass
601
+
602
+
603
+ def search_db(query: str, fields: List[str], where_clause: str = "", page: int = 1, results_per_page: int = 5) -> List[Dict[str, Any]]:
604
+ """
605
+ Perform a full-text search on specified fields with optional filtering and pagination.
606
+
607
+ Args:
608
+ query (str): The search query.
609
+ fields (List[str]): List of fields to search in.
610
+ where_clause (str, optional): Additional SQL WHERE clause to filter results.
611
+ page (int, optional): Page number for pagination.
612
+ results_per_page (int, optional): Number of results per page.
613
+
614
+ Returns:
615
+ List[Dict[str, Any]]: List of matching chat records with content and metadata.
616
+ """
617
+ if not query.strip():
618
+ return []
619
+
620
+ conn = sqlite3.connect(chat_DB_PATH)
621
+ cursor = conn.cursor()
622
+ try:
623
+ # Construct the MATCH query for FTS5
624
+ match_query = " AND ".join(fields) + f" MATCH ?"
625
+ # Adjust the query with the fields
626
+ fts_query = f"""
627
+ SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history
628
+ FROM CharacterChats_fts
629
+ JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id
630
+ WHERE {match_query}
631
+ """
632
+ if where_clause:
633
+ fts_query += f" AND ({where_clause})"
634
+ fts_query += " ORDER BY rank LIMIT ? OFFSET ?"
635
+ offset = (page - 1) * results_per_page
636
+ cursor.execute(fts_query, (query, results_per_page, offset))
637
+ rows = cursor.fetchall()
638
+ columns = [description[0] for description in cursor.description]
639
+ results = [dict(zip(columns, row)) for row in rows]
640
+ return results
641
+ except Exception as e:
642
+ logging.error(f"Error in search_db: {e}")
643
+ return []
644
+ finally:
645
+ conn.close()
646
+
647
+
648
+ def perform_full_text_search_chat(query: str, relevant_chat_ids: List[int], page: int = 1, results_per_page: int = 5) -> \
649
+ List[Dict[str, Any]]:
650
+ """
651
+ Perform a full-text search within the specified chat IDs using FTS5.
652
+
653
+ Args:
654
+ query (str): The user's query.
655
+ relevant_chat_ids (List[int]): List of chat IDs to search within.
656
+ page (int): Pagination page number.
657
+ results_per_page (int): Number of results per page.
658
+
659
+ Returns:
660
+ List[Dict[str, Any]]: List of search results with content and metadata.
661
+ """
662
+ try:
663
+ # Construct a WHERE clause to limit the search to relevant chat IDs
664
+ where_clause = " OR ".join([f"media_id = {chat_id}" for chat_id in relevant_chat_ids])
665
+ if not where_clause:
666
+ where_clause = "1" # No restriction if no chat IDs
667
+
668
+ # Perform full-text search using FTS5
669
+ fts_results = search_db(query, ["content"], where_clause, page=page, results_per_page=results_per_page)
670
+
671
+ filtered_fts_results = [
672
+ {
673
+ "content": result['content'],
674
+ "metadata": {"media_id": result['id']}
675
+ }
676
+ for result in fts_results
677
+ if result['id'] in relevant_chat_ids
678
+ ]
679
+ return filtered_fts_results
680
+ except Exception as e:
681
+ logging.error(f"Error in perform_full_text_search_chat: {str(e)}")
682
+ return []
683
+
684
+
685
+ def fetch_all_chats() -> List[Dict[str, Any]]:
686
+ """
687
+ Fetch all chat messages from the database.
688
+
689
+ Returns:
690
+ List[Dict[str, Any]]: List of chat messages with relevant metadata.
691
+ """
692
+ try:
693
+ chats = get_character_chats() # Modify this function to retrieve all chats
694
+ return chats
695
+ except Exception as e:
696
+ logging.error(f"Error fetching all chats: {str(e)}")
697
+ return []
698
+
699
+ #
700
+ # End of Character_Chat_DB.py
701
+ #######################################################################################################################
App_Function_Libraries/DB/DB_Manager.py ADDED
@@ -0,0 +1,991 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # DB_Manager.py
2
+ # Description: This file contains the DatabaseManager class, which is responsible for managing the database connection, i.e. either SQLite or Elasticsearch.
3
+ #
4
+ # Imports
5
+ import configparser
6
+ import os
7
+ import logging
8
+ import time
9
+ from typing import Tuple, List, Union, Dict
10
+ #
11
+ # 3rd-Party Libraries
12
+ from elasticsearch import Elasticsearch
13
+ #
14
+ # Import your existing SQLite functions
15
+ from App_Function_Libraries.DB.SQLite_DB import DatabaseError
16
+ from App_Function_Libraries.DB.SQLite_DB import (
17
+ update_media_content as sqlite_update_media_content,
18
+ list_prompts as sqlite_list_prompts,
19
+ search_and_display as sqlite_search_and_display,
20
+ fetch_prompt_details as sqlite_fetch_prompt_details,
21
+ keywords_browser_interface as sqlite_keywords_browser_interface,
22
+ add_keyword as sqlite_add_keyword,
23
+ delete_keyword as sqlite_delete_keyword,
24
+ export_keywords_to_csv as sqlite_export_keywords_to_csv,
25
+ ingest_article_to_db as sqlite_ingest_article_to_db,
26
+ add_media_to_database as sqlite_add_media_to_database,
27
+ import_obsidian_note_to_db as sqlite_import_obsidian_note_to_db,
28
+ add_prompt as sqlite_add_prompt,
29
+ delete_chat_message as sqlite_delete_chat_message,
30
+ update_chat_message as sqlite_update_chat_message,
31
+ add_chat_message as sqlite_add_chat_message,
32
+ get_chat_messages as sqlite_get_chat_messages,
33
+ search_chat_conversations as sqlite_search_chat_conversations,
34
+ create_chat_conversation as sqlite_create_chat_conversation,
35
+ save_chat_history_to_database as sqlite_save_chat_history_to_database,
36
+ view_database as sqlite_view_database,
37
+ get_transcripts as sqlite_get_transcripts,
38
+ get_trashed_items as sqlite_get_trashed_items,
39
+ user_delete_item as sqlite_user_delete_item,
40
+ empty_trash as sqlite_empty_trash,
41
+ create_automated_backup as sqlite_create_automated_backup,
42
+ add_or_update_prompt as sqlite_add_or_update_prompt,
43
+ load_prompt_details as sqlite_load_prompt_details,
44
+ load_preset_prompts as sqlite_load_preset_prompts,
45
+ insert_prompt_to_db as sqlite_insert_prompt_to_db,
46
+ delete_prompt as sqlite_delete_prompt,
47
+ search_and_display_items as sqlite_search_and_display_items,
48
+ get_conversation_name as sqlite_get_conversation_name,
49
+ add_media_with_keywords as sqlite_add_media_with_keywords,
50
+ check_media_and_whisper_model as sqlite_check_media_and_whisper_model, \
51
+ create_document_version as sqlite_create_document_version,
52
+ get_document_version as sqlite_get_document_version, sqlite_search_db, add_media_chunk as sqlite_add_media_chunk,
53
+ sqlite_update_fts_for_media, get_unprocessed_media as sqlite_get_unprocessed_media, fetch_item_details as sqlite_fetch_item_details, \
54
+ search_media_database as sqlite_search_media_database, mark_as_trash as sqlite_mark_as_trash, \
55
+ get_media_transcripts as sqlite_get_media_transcripts, get_specific_transcript as sqlite_get_specific_transcript, \
56
+ get_media_summaries as sqlite_get_media_summaries, get_specific_summary as sqlite_get_specific_summary, \
57
+ get_media_prompts as sqlite_get_media_prompts, get_specific_prompt as sqlite_get_specific_prompt, \
58
+ delete_specific_transcript as sqlite_delete_specific_transcript,
59
+ delete_specific_summary as sqlite_delete_specific_summary, \
60
+ delete_specific_prompt as sqlite_delete_specific_prompt,
61
+ fetch_keywords_for_media as sqlite_fetch_keywords_for_media, \
62
+ update_keywords_for_media as sqlite_update_keywords_for_media, check_media_exists as sqlite_check_media_exists, \
63
+ search_prompts as sqlite_search_prompts, get_media_content as sqlite_get_media_content, \
64
+ get_paginated_files as sqlite_get_paginated_files, get_media_title as sqlite_get_media_title, \
65
+ get_all_content_from_database as sqlite_get_all_content_from_database,
66
+ get_next_media_id as sqlite_get_next_media_id, \
67
+ batch_insert_chunks as sqlite_batch_insert_chunks, Database, save_workflow_chat_to_db as sqlite_save_workflow_chat_to_db, \
68
+ get_workflow_chat as sqlite_get_workflow_chat, update_media_content_with_version as sqlite_update_media_content_with_version, \
69
+ check_existing_media as sqlite_check_existing_media, get_all_document_versions as sqlite_get_all_document_versions, \
70
+ fetch_paginated_data as sqlite_fetch_paginated_data, get_latest_transcription as sqlite_get_latest_transcription, \
71
+ mark_media_as_processed as sqlite_mark_media_as_processed,
72
+ )
73
+ from App_Function_Libraries.DB.Character_Chat_DB import (
74
+ add_character_card as sqlite_add_character_card, get_character_cards as sqlite_get_character_cards, \
75
+ get_character_card_by_id as sqlite_get_character_card_by_id, update_character_card as sqlite_update_character_card, \
76
+ delete_character_card as sqlite_delete_character_card, add_character_chat as sqlite_add_character_chat, \
77
+ get_character_chats as sqlite_get_character_chats, get_character_chat_by_id as sqlite_get_character_chat_by_id, \
78
+ update_character_chat as sqlite_update_character_chat, delete_character_chat as sqlite_delete_character_chat, \
79
+ migrate_chat_to_media_db as sqlite_migrate_chat_to_media_db,
80
+ )
81
+ #
82
+ # Local Imports
83
+ from App_Function_Libraries.Utils.Utils import load_comprehensive_config, get_database_path, get_project_relative_path
84
+ #
85
+ # End of imports
86
+ ############################################################################################################
87
+
88
+
89
+ ############################################################################################################
90
+ #
91
+ # Database Config loading
92
+
93
+ logger = logging.getLogger(__name__)
94
+
95
+ config_path = get_project_relative_path('Config_Files/config.txt')
96
+ config = configparser.ConfigParser()
97
+ config.read(config_path)
98
+
99
+ db_path: str = config.get('Database', 'sqlite_path', fallback='./Databases/media_summary.db')
100
+ backup_path: str = config.get('Database', 'backup_path', fallback='database_backups')
101
+ backup_dir: Union[str, bytes] = os.environ.get('DB_BACKUP_DIR', backup_path)
102
+
103
+ def get_db_config():
104
+ try:
105
+ config = load_comprehensive_config()
106
+
107
+ if 'Database' not in config:
108
+ print("Warning: 'Database' section not found in config. Using default values.")
109
+ return default_db_config()
110
+
111
+ return {
112
+ 'type': config.get('Database', 'type', fallback='sqlite'),
113
+ 'sqlite_path': config.get('Database', 'sqlite_path', fallback='Databases/media_summary.db'),
114
+ 'elasticsearch_host': config.get('Database', 'elasticsearch_host', fallback='localhost'),
115
+ 'elasticsearch_port': config.getint('Database', 'elasticsearch_port', fallback=9200)
116
+ }
117
+ except FileNotFoundError:
118
+ print("Warning: Config file not found. Using default database configuration.")
119
+ return default_db_config()
120
+ except Exception as e:
121
+ print(f"Error reading config: {str(e)}. Using default database configuration.")
122
+ return default_db_config()
123
+
124
+ def default_db_config():
125
+ return {
126
+ 'type': 'sqlite',
127
+ 'sqlite_path': get_database_path('media_summary.db'),
128
+ 'elasticsearch_host': 'localhost',
129
+ 'elasticsearch_port': 9200
130
+ }
131
+
132
+ def ensure_directory_exists(file_path):
133
+ directory = os.path.dirname(file_path)
134
+ if not os.path.exists(directory):
135
+ os.makedirs(directory)
136
+ print(f"Created directory: {directory}")
137
+
138
+ db_config = get_db_config()
139
+ db_type = db_config['type']
140
+
141
+ if db_type == 'sqlite':
142
+ db = Database(os.path.basename(db_config['sqlite_path']))
143
+ elif db_type == 'elasticsearch':
144
+ raise NotImplementedError("Elasticsearch support not yet implemented")
145
+ else:
146
+ raise ValueError(f"Unsupported database type: {db_type}")
147
+
148
+ print(f"Database path: {db.db_path}")
149
+
150
+ def get_db_config():
151
+ try:
152
+ config = load_comprehensive_config()
153
+
154
+ if 'Database' not in config:
155
+ print("Warning: 'Database' section not found in config. Using default values.")
156
+ return default_db_config()
157
+
158
+ return {
159
+ 'type': config.get('Database', 'type', fallback='sqlite'),
160
+ 'sqlite_path': config.get('Database', 'sqlite_path', fallback='Databases/media_summary.db'),
161
+ 'elasticsearch_host': config.get('Database', 'elasticsearch_host', fallback='localhost'),
162
+ 'elasticsearch_port': config.getint('Database', 'elasticsearch_port', fallback=9200)
163
+ }
164
+ except FileNotFoundError:
165
+ print("Warning: Config file not found. Using default database configuration.")
166
+ return default_db_config()
167
+ except Exception as e:
168
+ print(f"Error reading config: {str(e)}. Using default database configuration.")
169
+ return default_db_config()
170
+
171
+
172
+ def default_db_config():
173
+ """Return the default database configuration with project-relative paths."""
174
+ return {
175
+ 'type': 'sqlite',
176
+ 'sqlite_path': get_database_path('media_summary.db'),
177
+ 'elasticsearch_host': 'localhost',
178
+ 'elasticsearch_port': 9200
179
+ }
180
+
181
+
182
+ def ensure_directory_exists(file_path):
183
+ directory = os.path.dirname(file_path)
184
+ if not os.path.exists(directory):
185
+ os.makedirs(directory)
186
+ print(f"Created directory: {directory}")
187
+
188
+ # Use the config to set up the database
189
+ db_config = get_db_config()
190
+ db_type = db_config['type']
191
+
192
+ if db_type == 'sqlite':
193
+ db = Database(os.path.basename(db_config['sqlite_path']))
194
+ elif db_type == 'elasticsearch':
195
+ # Implement Elasticsearch setup here if needed
196
+ raise NotImplementedError("Elasticsearch support not yet implemented")
197
+ else:
198
+ raise ValueError(f"Unsupported database type: {db_type}")
199
+
200
+ # Print database path for debugging
201
+ print(f"Database path: {db.db_path}")
202
+
203
+ # Sanity Check for SQLite DB
204
+ # FIXME - Remove this after testing / Writing Unit tests
205
+ # try:
206
+ # db.execute_query("CREATE TABLE IF NOT EXISTS test_table (id INTEGER PRIMARY KEY)")
207
+ # logger.info("Successfully created test table")
208
+ # except DatabaseError as e:
209
+ # logger.error(f"Failed to create test table: {e}")
210
+
211
+ #
212
+ # End of Database Config loading
213
+ ############################################################################################################
214
+ #
215
+ # DB Search functions
216
+
217
+ def search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10):
218
+ if db_type == 'sqlite':
219
+ return sqlite_search_db(search_query, search_fields, keywords, page, results_per_page)
220
+ elif db_type == 'elasticsearch':
221
+ # Implement Elasticsearch version when available
222
+ raise NotImplementedError("Elasticsearch version of search_db not yet implemented")
223
+ else:
224
+ raise ValueError(f"Unsupported database type: {db_type}")
225
+
226
+ def view_database(*args, **kwargs):
227
+ if db_type == 'sqlite':
228
+ return sqlite_view_database(*args, **kwargs)
229
+ elif db_type == 'elasticsearch':
230
+ # Implement Elasticsearch version
231
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
232
+
233
+ def search_and_display_items(*args, **kwargs):
234
+ if db_type == 'sqlite':
235
+ return sqlite_search_and_display_items(*args, **kwargs)
236
+ elif db_type == 'elasticsearch':
237
+ # Implement Elasticsearch version
238
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
239
+
240
+ def get_all_content_from_database():
241
+ if db_type == 'sqlite':
242
+ return sqlite_get_all_content_from_database()
243
+ elif db_type == 'elasticsearch':
244
+ # Implement Elasticsearch version
245
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
246
+
247
+ def search_and_display(*args, **kwargs):
248
+ if db_type == 'sqlite':
249
+ return sqlite_search_and_display(*args, **kwargs)
250
+ elif db_type == 'elasticsearch':
251
+ # Implement Elasticsearch version
252
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
253
+
254
+ def check_media_exists(*args, **kwargs):
255
+ if db_type == 'sqlite':
256
+ return sqlite_check_media_exists(*args, **kwargs)
257
+ elif db_type == 'elasticsearch':
258
+ # Implement Elasticsearch version
259
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
260
+
261
+ def get_paginated_files(*args, **kwargs):
262
+ if db_type == 'sqlite':
263
+ return sqlite_get_paginated_files(*args, **kwargs)
264
+ elif db_type == 'elasticsearch':
265
+ # Implement Elasticsearch version
266
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
267
+
268
+ def get_media_title(*args, **kwargs):
269
+ if db_type == 'sqlite':
270
+ return sqlite_get_media_title(*args, **kwargs)
271
+ elif db_type == 'elasticsearch':
272
+ # Implement Elasticsearch version
273
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
274
+
275
+ def get_next_media_id():
276
+ if db_type == 'sqlite':
277
+ return sqlite_get_next_media_id()
278
+ elif db_type == 'elasticsearch':
279
+ # Implement Elasticsearch version
280
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
281
+
282
+ #
283
+ # End of DB-Searching functions
284
+ ############################################################################################################
285
+
286
+
287
+ ############################################################################################################
288
+ #
289
+ # Transcript-related Functions
290
+
291
+ def get_transcripts(*args, **kwargs):
292
+ if db_type == 'sqlite':
293
+ return sqlite_get_transcripts(*args, **kwargs)
294
+ elif db_type == 'elasticsearch':
295
+ # Implement Elasticsearch version
296
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
297
+
298
+ #
299
+ # End of Transcript-related Functions
300
+ ############################################################################################################
301
+
302
+
303
+ ############################################################################################################
304
+ #
305
+ # DB-Ingestion functions
306
+
307
+ def add_media_to_database(*args, **kwargs):
308
+ if db_type == 'sqlite':
309
+ result = sqlite_add_media_to_database(*args, **kwargs)
310
+
311
+ # Extract content
312
+ segments = kwargs.get('segments') if 'segments' in kwargs else args[2] if len(args) > 2 else None
313
+ if segments is None:
314
+ raise ValueError("Segments not provided in arguments")
315
+
316
+ if isinstance(segments, list):
317
+ content = ' '.join([segment.get('Text', '') for segment in segments if 'Text' in segment])
318
+ elif isinstance(segments, dict):
319
+ content = segments.get('text', '') or segments.get('content', '')
320
+ else:
321
+ content = str(segments)
322
+
323
+ # Extract media_id from the result
324
+ # Assuming the result is in the format "Media 'Title' added/updated successfully with ID: {media_id}"
325
+ import re
326
+ match = re.search(r"with ID: (\d+)", result)
327
+ if match:
328
+ media_id = int(match.group(1))
329
+
330
+ # Create initial document version
331
+ sqlite_create_document_version(media_id, content)
332
+
333
+ return result
334
+ elif db_type == 'elasticsearch':
335
+ # Implement Elasticsearch version
336
+ raise NotImplementedError("Elasticsearch version of add_media_to_database not yet implemented")
337
+
338
+ def check_existing_media(*args, **kwargs):
339
+ if db_type == 'sqlite':
340
+ return sqlite_check_existing_media(*args, **kwargs)
341
+ elif db_type == 'elasticsearch':
342
+ # Implement Elasticsearch version
343
+ raise NotImplementedError("Elasticsearch version of check_existing_media not yet implemented")
344
+
345
+ def update_media_content_with_version(*args, **kwargs):
346
+ if db_type == 'sqlite':
347
+ return sqlite_update_media_content_with_version(*args, **kwargs)
348
+ elif db_type == 'elasticsearch':
349
+ # Implement Elasticsearch version
350
+ raise NotImplementedError("Elasticsearch version of update_media_content not yet implemented")
351
+
352
+ def import_obsidian_note_to_db(*args, **kwargs):
353
+ if db_type == 'sqlite':
354
+ return sqlite_import_obsidian_note_to_db(*args, **kwargs)
355
+ elif db_type == 'elasticsearch':
356
+ # Implement Elasticsearch version
357
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
358
+
359
+
360
+ def update_media_content(*args, **kwargs):
361
+ if db_type == 'sqlite':
362
+ result = sqlite_update_media_content(*args, **kwargs)
363
+
364
+ # Extract media_id and content
365
+ selected_item = args[0]
366
+ item_mapping = args[1]
367
+ content_input = args[2]
368
+
369
+ if selected_item and item_mapping and selected_item in item_mapping:
370
+ media_id = item_mapping[selected_item]
371
+
372
+ # Create new document version
373
+ sqlite_create_document_version(media_id, content_input)
374
+
375
+ return result
376
+ elif db_type == 'elasticsearch':
377
+ # Implement Elasticsearch version
378
+ raise NotImplementedError("Elasticsearch version of update_media_content not yet implemented")
379
+
380
+
381
+ def add_media_with_keywords(*args, **kwargs):
382
+ if db_type == 'sqlite':
383
+ return sqlite_add_media_with_keywords(*args, **kwargs)
384
+ elif db_type == 'elasticsearch':
385
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
386
+
387
+ def check_media_and_whisper_model(*args, **kwargs):
388
+ if db_type == 'sqlite':
389
+ return sqlite_check_media_and_whisper_model(*args, **kwargs)
390
+ elif db_type == 'elasticsearch':
391
+ raise NotImplementedError("Elasticsearch version of check_media_and_whisper_model not yet implemented")
392
+
393
+ def ingest_article_to_db(url, title, author, content, keywords, summary, ingestion_date, custom_prompt):
394
+ if db_type == 'sqlite':
395
+ return sqlite_ingest_article_to_db(url, title, author, content, keywords, summary, ingestion_date, custom_prompt)
396
+ elif db_type == 'elasticsearch':
397
+ # Implement Elasticsearch version
398
+ raise NotImplementedError("Elasticsearch version of ingest_article_to_db not yet implemented")
399
+ else:
400
+ raise ValueError(f"Unsupported database type: {db_type}")
401
+
402
+
403
+ def add_media_chunk(*args, **kwargs):
404
+ if db_type == 'sqlite':
405
+ sqlite_add_media_chunk(*args, **kwargs)
406
+ elif db_type == 'elasticsearch':
407
+ # Implement Elasticsearch version
408
+ raise NotImplementedError("Elasticsearch version not yet implemented")
409
+ else:
410
+ raise ValueError(f"Unsupported database type: {db_type}")
411
+
412
+ def batch_insert_chunks(*args, **kwargs):
413
+ if db_type == 'sqlite':
414
+ sqlite_batch_insert_chunks(*args, **kwargs)
415
+ elif db_type == 'elasticsearch':
416
+ # Implement Elasticsearch version
417
+ raise NotImplementedError("Elasticsearch version not yet implemented")
418
+ else:
419
+ raise ValueError(f"Unsupported database type: {db_type}")
420
+
421
+ def update_fts_for_media(media_id: int):
422
+ if db_type == 'sqlite':
423
+ sqlite_update_fts_for_media(db, media_id)
424
+ elif db_type == 'elasticsearch':
425
+ # Implement Elasticsearch version
426
+ raise NotImplementedError("Elasticsearch version not yet implemented")
427
+ else:
428
+ raise ValueError(f"Unsupported database type: {db_type}")
429
+
430
+
431
+ def get_unprocessed_media(*args, **kwargs):
432
+ if db_type == 'sqlite':
433
+ return sqlite_get_unprocessed_media(db)
434
+ elif db_type == 'elasticsearch':
435
+ # Implement Elasticsearch version
436
+ raise NotImplementedError("Elasticsearch version of get_unprocessed_media not yet implemented")
437
+ else:
438
+ raise ValueError(f"Unsupported database type: {db_type}")
439
+
440
+
441
+ def mark_media_as_processed(*args, **kwargs):
442
+ if db_type == 'sqlite':
443
+ return sqlite_mark_media_as_processed(*args, **kwargs)
444
+ elif db_type == 'elasticsearch':
445
+ # Implement Elasticsearch version
446
+ raise NotImplementedError("Elasticsearch version of mark_media_as_processed not yet implemented")
447
+ else:
448
+ raise ValueError(f"Unsupported database type: {db_type}")
449
+
450
+
451
+ #
452
+ # End of DB-Ingestion functions
453
+ ############################################################################################################
454
+
455
+
456
+ ############################################################################################################
457
+ #
458
+ # Prompt-related functions #FIXME rename /resort
459
+
460
+ def list_prompts(*args, **kwargs):
461
+ if db_type == 'sqlite':
462
+ return sqlite_list_prompts(*args, **kwargs)
463
+ elif db_type == 'elasticsearch':
464
+ # Implement Elasticsearch version
465
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
466
+
467
+ def search_prompts(query):
468
+ if db_type == 'sqlite':
469
+ return sqlite_search_prompts(query)
470
+ elif db_type == 'elasticsearch':
471
+ # Implement Elasticsearch version
472
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
473
+
474
+ def fetch_prompt_details(*args, **kwargs):
475
+ if db_type == 'sqlite':
476
+ return sqlite_fetch_prompt_details(*args, **kwargs)
477
+ elif db_type == 'elasticsearch':
478
+ # Implement Elasticsearch version
479
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
480
+
481
+ def add_prompt(*args, **kwargs):
482
+ if db_type == 'sqlite':
483
+ return sqlite_add_prompt(*args, **kwargs)
484
+ elif db_type == 'elasticsearch':
485
+ # Implement Elasticsearch version
486
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
487
+
488
+
489
+ def add_or_update_prompt(*args, **kwargs):
490
+ if db_type == 'sqlite':
491
+ return sqlite_add_or_update_prompt(*args, **kwargs)
492
+ elif db_type == 'elasticsearch':
493
+ # Implement Elasticsearch version
494
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
495
+
496
+ def load_prompt_details(*args, **kwargs):
497
+ if db_type == 'sqlite':
498
+ return sqlite_load_prompt_details(*args, **kwargs)
499
+ elif db_type == 'elasticsearch':
500
+ # Implement Elasticsearch version
501
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
502
+
503
+ def load_preset_prompts(*args, **kwargs):
504
+ if db_type == 'sqlite':
505
+ return sqlite_load_preset_prompts()
506
+ elif db_type == 'elasticsearch':
507
+ # Implement Elasticsearch version
508
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
509
+
510
+ def insert_prompt_to_db(*args, **kwargs):
511
+ if db_type == 'sqlite':
512
+ return sqlite_insert_prompt_to_db(*args, **kwargs)
513
+ elif db_type == 'elasticsearch':
514
+ # Implement Elasticsearch version
515
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
516
+
517
+ def delete_prompt(*args, **kwargs):
518
+ if db_type == 'sqlite':
519
+ return sqlite_delete_prompt(*args, **kwargs)
520
+ elif db_type == 'elasticsearch':
521
+ # Implement Elasticsearch version
522
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
523
+
524
+ def search_media_database(*args, **kwargs):
525
+ if db_type == 'sqlite':
526
+ return sqlite_search_media_database(*args, **kwargs)
527
+ elif db_type == 'elasticsearch':
528
+ # Implement Elasticsearch version when available
529
+ raise NotImplementedError("Elasticsearch version of search_media_database not yet implemented")
530
+ else:
531
+ raise ValueError(f"Unsupported database type: {db_type}")
532
+
533
+ def mark_as_trash(media_id: int) -> None:
534
+ if db_type == 'sqlite':
535
+ return sqlite_mark_as_trash(media_id)
536
+ elif db_type == 'elasticsearch':
537
+ # Implement Elasticsearch version when available
538
+ raise NotImplementedError("Elasticsearch version of mark_as_trash not yet implemented")
539
+ else:
540
+ raise ValueError(f"Unsupported database type: {db_type}")
541
+
542
+
543
+ def get_latest_transcription(*args, **kwargs):
544
+ if db_type == 'sqlite':
545
+ return sqlite_get_latest_transcription(*args, **kwargs)
546
+ elif db_type == 'elasticsearch':
547
+ # Implement Elasticsearch version
548
+ raise NotImplementedError("Elasticsearch version of get_latest_transcription not yet implemented")
549
+
550
+ def fetch_paginated_data(*args, **kwargs):
551
+ if db_type == 'sqlite':
552
+ return sqlite_fetch_paginated_data(*args, **kwargs)
553
+ elif db_type == 'elasticsearch':
554
+ # Implement Elasticsearch version
555
+ raise NotImplementedError("Elasticsearch version of fetch_paginated_data not yet implemented")
556
+ else:
557
+ raise ValueError(f"Unsupported database type: {db_type}")
558
+
559
+
560
+ def get_media_content(media_id: int) -> str:
561
+ if db_type == 'sqlite':
562
+ return sqlite_get_media_content(media_id)
563
+ elif db_type == 'elasticsearch':
564
+ raise NotImplementedError("Elasticsearch version of get_media_content not yet implemented")
565
+ else:
566
+ raise ValueError(f"Unsupported database type: {db_type}")
567
+
568
+ def get_media_transcripts(media_id: int) -> List[Dict]:
569
+ if db_type == 'sqlite':
570
+ return sqlite_get_media_transcripts(media_id)
571
+ elif db_type == 'elasticsearch':
572
+ raise NotImplementedError("Elasticsearch version of get_media_transcripts not yet implemented")
573
+ else:
574
+ raise ValueError(f"Unsupported database type: {db_type}")
575
+
576
+ def get_specific_transcript(transcript_id: int) -> Dict:
577
+ if db_type == 'sqlite':
578
+ return sqlite_get_specific_transcript(transcript_id)
579
+ elif db_type == 'elasticsearch':
580
+ raise NotImplementedError("Elasticsearch version of get_specific_transcript not yet implemented")
581
+ else:
582
+ raise ValueError(f"Unsupported database type: {db_type}")
583
+
584
+ def get_media_summaries(media_id: int) -> List[Dict]:
585
+ if db_type == 'sqlite':
586
+ return sqlite_get_media_summaries(media_id)
587
+ elif db_type == 'elasticsearch':
588
+ raise NotImplementedError("Elasticsearch version of get_media_summaries not yet implemented")
589
+ else:
590
+ raise ValueError(f"Unsupported database type: {db_type}")
591
+
592
+ def get_specific_summary(summary_id: int) -> Dict:
593
+ if db_type == 'sqlite':
594
+ return sqlite_get_specific_summary(summary_id)
595
+ elif db_type == 'elasticsearch':
596
+ raise NotImplementedError("Elasticsearch version of get_specific_summary not yet implemented")
597
+ else:
598
+ raise ValueError(f"Unsupported database type: {db_type}")
599
+
600
+ def fetch_item_details_single(*args, **kwargs):
601
+ if db_type == 'sqlite':
602
+ return sqlite_fetch_item_details(*args, **kwargs)
603
+ elif db_type == 'elasticsearch':
604
+ # Implement Elasticsearch version
605
+ raise NotImplementedError("Elasticsearch version of fetch_item_details not yet implemented")
606
+ else:
607
+ raise ValueError(f"Unsupported database type: {db_type}")
608
+
609
+ def get_all_document_versions(*args, **kwargs):
610
+ if db_type == 'sqlite':
611
+ return sqlite_get_all_document_versions(*args, **kwargs)
612
+ elif db_type == 'elasticsearch':
613
+ # Implement Elasticsearch version
614
+ raise NotImplementedError("Elasticsearch version of get_all_document_versions not yet implemented")
615
+ else:
616
+ raise ValueError(f"Unsupported database type: {db_type}")
617
+ #
618
+ #
619
+ ############################################################################################################
620
+ #
621
+ # Prompt Functions:
622
+
623
+ def get_media_prompts(media_id: int) -> List[Dict]:
624
+ if db_type == 'sqlite':
625
+ return sqlite_get_media_prompts(media_id)
626
+ elif db_type == 'elasticsearch':
627
+ raise NotImplementedError("Elasticsearch version of get_media_prompts not yet implemented")
628
+ else:
629
+ raise ValueError(f"Unsupported database type: {db_type}")
630
+
631
+ def get_specific_prompt(prompt_id: int) -> Dict:
632
+ if db_type == 'sqlite':
633
+ return sqlite_get_specific_prompt(prompt_id)
634
+ elif db_type == 'elasticsearch':
635
+ raise NotImplementedError("Elasticsearch version of get_specific_prompt not yet implemented")
636
+ else:
637
+ return {'error': f"Unsupported database type: {db_type}"}
638
+
639
+ def delete_specific_transcript(transcript_id: int) -> str:
640
+ if db_type == 'sqlite':
641
+ return sqlite_delete_specific_transcript(transcript_id)
642
+ elif db_type == 'elasticsearch':
643
+ raise NotImplementedError("Elasticsearch version of delete_specific_transcript not yet implemented")
644
+ else:
645
+ raise ValueError(f"Unsupported database type: {db_type}")
646
+
647
+ def delete_specific_summary(summary_id: int) -> str:
648
+ if db_type == 'sqlite':
649
+ return sqlite_delete_specific_summary(summary_id)
650
+ elif db_type == 'elasticsearch':
651
+ raise NotImplementedError("Elasticsearch version of delete_specific_summary not yet implemented")
652
+ else:
653
+ raise ValueError(f"Unsupported database type: {db_type}")
654
+
655
+ def delete_specific_prompt(prompt_id: int) -> str:
656
+ if db_type == 'sqlite':
657
+ return sqlite_delete_specific_prompt(prompt_id)
658
+ elif db_type == 'elasticsearch':
659
+ raise NotImplementedError("Elasticsearch version of delete_specific_prompt not yet implemented")
660
+ else:
661
+ raise ValueError(f"Unsupported database type: {db_type}")
662
+
663
+
664
+ #
665
+ # End of Prompt-related functions
666
+ ############################################################################################################
667
+
668
+ ############################################################################################################
669
+ #
670
+ # Keywords-related Functions
671
+
672
+ def keywords_browser_interface(*args, **kwargs):
673
+ if db_type == 'sqlite':
674
+ return sqlite_keywords_browser_interface()
675
+ elif db_type == 'elasticsearch':
676
+ # Implement Elasticsearch version
677
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
678
+
679
+ def add_keyword(*args, **kwargs):
680
+ if db_type == 'sqlite':
681
+ with db.get_connection() as conn:
682
+ cursor = conn.cursor()
683
+ return sqlite_add_keyword(*args, **kwargs)
684
+ elif db_type == 'elasticsearch':
685
+ # Implement Elasticsearch version
686
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
687
+
688
+ def delete_keyword(*args, **kwargs):
689
+ if db_type == 'sqlite':
690
+ return sqlite_delete_keyword(*args, **kwargs)
691
+ elif db_type == 'elasticsearch':
692
+ # Implement Elasticsearch version
693
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
694
+
695
+ def export_keywords_to_csv(*args, **kwargs):
696
+ if db_type == 'sqlite':
697
+ return sqlite_export_keywords_to_csv()
698
+ elif db_type == 'elasticsearch':
699
+ # Implement Elasticsearch version
700
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
701
+
702
+ def update_keywords_for_media(*args, **kwargs):
703
+ if db_type == 'sqlite':
704
+ return sqlite_update_keywords_for_media(*args, **kwargs)
705
+ elif db_type == 'elasticsearch':
706
+ # Implement Elasticsearch version
707
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
708
+
709
+ def fetch_keywords_for_media(*args, **kwargs):
710
+ if db_type == 'sqlite':
711
+ return sqlite_fetch_keywords_for_media(*args, **kwargs)
712
+ elif db_type == 'elasticsearch':
713
+ # Implement Elasticsearch version
714
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
715
+
716
+ #
717
+ # End of Keywords-related Functions
718
+ ############################################################################################################
719
+
720
+ ############################################################################################################
721
+ #
722
+ # Chat-related Functions
723
+
724
+ def delete_chat_message(*args, **kwargs):
725
+ if db_type == 'sqlite':
726
+ return sqlite_delete_chat_message(*args, **kwargs)
727
+ elif db_type == 'elasticsearch':
728
+ # Implement Elasticsearch version
729
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
730
+
731
+ def update_chat_message(*args, **kwargs):
732
+ if db_type == 'sqlite':
733
+ return sqlite_update_chat_message(*args, **kwargs)
734
+ elif db_type == 'elasticsearch':
735
+ # Implement Elasticsearch version
736
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
737
+
738
+ def add_chat_message(*args, **kwargs):
739
+ if db_type == 'sqlite':
740
+ return sqlite_add_chat_message(*args, **kwargs)
741
+ elif db_type == 'elasticsearch':
742
+ # Implement Elasticsearch version
743
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
744
+
745
+ def get_chat_messages(*args, **kwargs):
746
+ if db_type == 'sqlite':
747
+ return sqlite_get_chat_messages(*args, **kwargs)
748
+ elif db_type == 'elasticsearch':
749
+ # Implement Elasticsearch version
750
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
751
+
752
+ def search_chat_conversations(*args, **kwargs):
753
+ if db_type == 'sqlite':
754
+ return sqlite_search_chat_conversations(*args, **kwargs)
755
+ elif db_type == 'elasticsearch':
756
+ # Implement Elasticsearch version
757
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
758
+
759
+ def create_chat_conversation(*args, **kwargs):
760
+ if db_type == 'sqlite':
761
+ return sqlite_create_chat_conversation(*args, **kwargs)
762
+ elif db_type == 'elasticsearch':
763
+ # Implement Elasticsearch version
764
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
765
+
766
+ def save_chat_history_to_database(*args, **kwargs):
767
+ if db_type == 'sqlite':
768
+ return sqlite_save_chat_history_to_database(*args, **kwargs)
769
+ elif db_type == 'elasticsearch':
770
+ # Implement Elasticsearch version
771
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
772
+
773
+ def get_conversation_name(*args, **kwargs):
774
+ if db_type == 'sqlite':
775
+ return sqlite_get_conversation_name(*args, **kwargs)
776
+ elif db_type == 'elasticsearch':
777
+ # Implement Elasticsearch version
778
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
779
+
780
+ #
781
+ # End of Chat-related Functions
782
+ ############################################################################################################
783
+
784
+
785
+ ############################################################################################################
786
+ #
787
+ # Character Chat-related Functions
788
+
789
+ def add_character_card(*args, **kwargs):
790
+ if db_type == 'sqlite':
791
+ return sqlite_add_character_card(*args, **kwargs)
792
+ elif db_type == 'elasticsearch':
793
+ # Implement Elasticsearch version
794
+ raise NotImplementedError("Elasticsearch version of add_character_card not yet implemented")
795
+
796
+ def get_character_cards():
797
+ if db_type == 'sqlite':
798
+ return sqlite_get_character_cards()
799
+ elif db_type == 'elasticsearch':
800
+ # Implement Elasticsearch version
801
+ raise NotImplementedError("Elasticsearch version of get_character_cards not yet implemented")
802
+
803
+ def get_character_card_by_id(*args, **kwargs):
804
+ if db_type == 'sqlite':
805
+ return sqlite_get_character_card_by_id(*args, **kwargs)
806
+ elif db_type == 'elasticsearch':
807
+ # Implement Elasticsearch version
808
+ raise NotImplementedError("Elasticsearch version of get_character_card_by_id not yet implemented")
809
+
810
+ def update_character_card(*args, **kwargs):
811
+ if db_type == 'sqlite':
812
+ return sqlite_update_character_card(*args, **kwargs)
813
+ elif db_type == 'elasticsearch':
814
+ # Implement Elasticsearch version
815
+ raise NotImplementedError("Elasticsearch version of update_character_card not yet implemented")
816
+
817
+ def delete_character_card(*args, **kwargs):
818
+ if db_type == 'sqlite':
819
+ return sqlite_delete_character_card(*args, **kwargs)
820
+ elif db_type == 'elasticsearch':
821
+ # Implement Elasticsearch version
822
+ raise NotImplementedError("Elasticsearch version of delete_character_card not yet implemented")
823
+
824
+ def add_character_chat(*args, **kwargs):
825
+ if db_type == 'sqlite':
826
+ return sqlite_add_character_chat(*args, **kwargs)
827
+ elif db_type == 'elasticsearch':
828
+ # Implement Elasticsearch version
829
+ raise NotImplementedError("Elasticsearch version of add_character_chat not yet implemented")
830
+
831
+ def get_character_chats(*args, **kwargs):
832
+ if db_type == 'sqlite':
833
+ return sqlite_get_character_chats(*args, **kwargs)
834
+ elif db_type == 'elasticsearch':
835
+ # Implement Elasticsearch version
836
+ raise NotImplementedError("Elasticsearch version of get_character_chats not yet implemented")
837
+
838
+ def get_character_chat_by_id(*args, **kwargs):
839
+ if db_type == 'sqlite':
840
+ return sqlite_get_character_chat_by_id(*args, **kwargs)
841
+ elif db_type == 'elasticsearch':
842
+ # Implement Elasticsearch version
843
+ raise NotImplementedError("Elasticsearch version of get_character_chat_by_id not yet implemented")
844
+
845
+ def update_character_chat(*args, **kwargs):
846
+ if db_type == 'sqlite':
847
+ return sqlite_update_character_chat(*args, **kwargs)
848
+ elif db_type == 'elasticsearch':
849
+ # Implement Elasticsearch version
850
+ raise NotImplementedError("Elasticsearch version of update_character_chat not yet implemented")
851
+
852
+ def delete_character_chat(*args, **kwargs):
853
+ if db_type == 'sqlite':
854
+ return sqlite_delete_character_chat(*args, **kwargs)
855
+ elif db_type == 'elasticsearch':
856
+ # Implement Elasticsearch version
857
+ raise NotImplementedError("Elasticsearch version of delete_character_chat not yet implemented")
858
+
859
+ def migrate_chat_to_media_db(*args, **kwargs):
860
+ if db_type == 'sqlite':
861
+ return sqlite_migrate_chat_to_media_db(*args, **kwargs)
862
+ elif db_type == 'elasticsearch':
863
+ # Implement Elasticsearch version
864
+ raise NotImplementedError("Elasticsearch version of migrate_chat_to_media_db not yet implemented")
865
+
866
+ #
867
+ # End of Character Chat-related Functions
868
+ ############################################################################################################
869
+
870
+
871
+ ############################################################################################################
872
+ #
873
+ # Trash-related Functions
874
+
875
+ def get_trashed_items(*args, **kwargs):
876
+ if db_type == 'sqlite':
877
+ return sqlite_get_trashed_items()
878
+ elif db_type == 'elasticsearch':
879
+ # Implement Elasticsearch version
880
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
881
+
882
+ def user_delete_item(*args, **kwargs):
883
+ if db_type == 'sqlite':
884
+ return sqlite_user_delete_item(*args, **kwargs)
885
+ elif db_type == 'elasticsearch':
886
+ # Implement Elasticsearch version
887
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
888
+
889
+ def empty_trash(*args, **kwargs):
890
+ if db_type == 'sqlite':
891
+ return sqlite_empty_trash(*args, **kwargs)
892
+ elif db_type == 'elasticsearch':
893
+ # Implement Elasticsearch version
894
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
895
+
896
+
897
+ def fetch_item_details(media_id: int) -> Tuple[str, str, str]:
898
+ """
899
+ Fetch the details of a media item including content, prompt, and summary.
900
+
901
+ Args:
902
+ media_id (int): The ID of the media item.
903
+
904
+ Returns:
905
+ Tuple[str, str, str]: A tuple containing (content, prompt, summary).
906
+ If an error occurs, it returns empty strings for each field.
907
+ """
908
+ if db_type == 'sqlite':
909
+ return sqlite_fetch_item_details(media_id)
910
+ elif db_type == 'elasticsearch':
911
+ # Implement Elasticsearch version when available
912
+ raise NotImplementedError("Elasticsearch version of fetch_item_details not yet implemented")
913
+ else:
914
+ raise ValueError(f"Unsupported database type: {db_type}")
915
+
916
+ #
917
+ # End of Trash-related Functions
918
+ ############################################################################################################
919
+
920
+
921
+ ############################################################################################################
922
+ #
923
+ # DB-Backup Functions
924
+
925
+ def create_automated_backup(*args, **kwargs):
926
+ if db_type == 'sqlite':
927
+ return sqlite_create_automated_backup(*args, **kwargs)
928
+ elif db_type == 'elasticsearch':
929
+ # Implement Elasticsearch version
930
+ raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
931
+
932
+ #
933
+ # End of DB-Backup Functions
934
+ ############################################################################################################
935
+
936
+
937
+ ############################################################################################################
938
+ #
939
+ # Document Versioning Functions
940
+
941
+ def create_document_version(*args, **kwargs):
942
+ if db_type == 'sqlite':
943
+ return sqlite_create_document_version(*args, **kwargs)
944
+ elif db_type == 'elasticsearch':
945
+ # Implement Elasticsearch version
946
+ raise NotImplementedError("Elasticsearch version of create_document_version not yet implemented")
947
+
948
+ def get_document_version(*args, **kwargs):
949
+ if db_type == 'sqlite':
950
+ return sqlite_get_document_version(*args, **kwargs)
951
+ elif db_type == 'elasticsearch':
952
+ # Implement Elasticsearch version
953
+ raise NotImplementedError("Elasticsearch version of get_document_version not yet implemented")
954
+
955
+ #
956
+ # End of Document Versioning Functions
957
+ ############################################################################################################
958
+
959
+
960
+ ############################################################################################################
961
+ #
962
+ # Workflow Functions
963
+
964
+ def get_workflow_chat(*args, **kwargs):
965
+ if db_type == 'sqlite':
966
+ return sqlite_get_workflow_chat(*args, **kwargs)
967
+ elif db_type == 'elasticsearch':
968
+ # Implement Elasticsearch version
969
+ raise NotImplementedError("Elasticsearch version of get_workflow_chat not yet implemented")
970
+
971
+
972
+ def save_workflow_chat_to_db(*args, **kwargs):
973
+ if db_type == 'sqlite':
974
+ # FIXME
975
+ return sqlite_save_workflow_chat_to_db(*args, **kwargs)
976
+ elif db_type == 'elasticsearch':
977
+ # Implement Elasticsearch version
978
+ raise NotImplementedError("Elasticsearch version of save_workflow_chat_to_db not yet implemented")
979
+
980
+ #
981
+ # End of Workflow Functions
982
+ ############################################################################################################
983
+
984
+ # Dead code FIXME
985
+ # def close_connection():
986
+ # if db_type == 'sqlite':
987
+ # db.get_connection().close()
988
+
989
+ #
990
+ # End of file
991
+ ############################################################################################################
App_Function_Libraries/DB/RAG_QA_Chat_DB.py ADDED
@@ -0,0 +1,722 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # RAG_QA_Chat_DB.py
2
+ # Description: This file contains the database operations for the RAG QA Chat + Notes system.
3
+ #
4
+ # Imports
5
+ import configparser
6
+ import logging
7
+ import re
8
+ import sqlite3
9
+ import uuid
10
+ from contextlib import contextmanager
11
+ from datetime import datetime
12
+
13
+ from App_Function_Libraries.Utils.Utils import get_project_relative_path, get_database_path
14
+
15
+ #
16
+ # External Imports
17
+ # (No external imports)
18
+ #
19
+ # Local Imports
20
+ # (No additional local imports)
21
+ #
22
+ ########################################################################################################################
23
+ #
24
+ # Functions:
25
+
26
+ # Construct the path to the config file
27
+ config_path = get_project_relative_path('Config_Files/config.txt')
28
+
29
+ # Read the config file
30
+ config = configparser.ConfigParser()
31
+ config.read(config_path)
32
+
33
+ # Get the SQLite path from the config, or use the default if not specified
34
+ if config.has_section('Database') and config.has_option('Database', 'rag_qa_db_path'):
35
+ rag_qa_db_path = config.get('Database', 'rag_qa_db_path')
36
+ else:
37
+ rag_qa_db_path = get_database_path('RAG_QA_Chat.db')
38
+
39
+ print(f"RAG QA Chat Database path: {rag_qa_db_path}")
40
+
41
+ # Set up logging
42
+ logging.basicConfig(level=logging.INFO)
43
+ logger = logging.getLogger(__name__)
44
+
45
+ # Database schema
46
+ SCHEMA_SQL = '''
47
+ -- Table for storing chat messages
48
+ CREATE TABLE IF NOT EXISTS rag_qa_chats (
49
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
50
+ conversation_id TEXT NOT NULL,
51
+ timestamp DATETIME NOT NULL,
52
+ role TEXT NOT NULL,
53
+ content TEXT NOT NULL
54
+ );
55
+
56
+ -- Table for storing conversation metadata
57
+ CREATE TABLE IF NOT EXISTS conversation_metadata (
58
+ conversation_id TEXT PRIMARY KEY,
59
+ created_at DATETIME NOT NULL,
60
+ last_updated DATETIME NOT NULL,
61
+ title TEXT NOT NULL
62
+ );
63
+
64
+ -- Table for storing keywords
65
+ CREATE TABLE IF NOT EXISTS rag_qa_keywords (
66
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
67
+ keyword TEXT NOT NULL UNIQUE
68
+ );
69
+
70
+ -- Table for linking keywords to conversations
71
+ CREATE TABLE IF NOT EXISTS rag_qa_conversation_keywords (
72
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
73
+ conversation_id TEXT NOT NULL,
74
+ keyword_id INTEGER NOT NULL,
75
+ FOREIGN KEY (conversation_id) REFERENCES conversation_metadata(conversation_id),
76
+ FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id)
77
+ );
78
+
79
+ -- Table for storing keyword collections
80
+ CREATE TABLE IF NOT EXISTS rag_qa_keyword_collections (
81
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
82
+ name TEXT NOT NULL UNIQUE,
83
+ parent_id INTEGER,
84
+ FOREIGN KEY (parent_id) REFERENCES rag_qa_keyword_collections(id)
85
+ );
86
+
87
+ -- Table for linking keywords to collections
88
+ CREATE TABLE IF NOT EXISTS rag_qa_collection_keywords (
89
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
90
+ collection_id INTEGER NOT NULL,
91
+ keyword_id INTEGER NOT NULL,
92
+ FOREIGN KEY (collection_id) REFERENCES rag_qa_keyword_collections(id),
93
+ FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id)
94
+ );
95
+
96
+ -- Table for storing notes
97
+ CREATE TABLE IF NOT EXISTS rag_qa_notes (
98
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
99
+ conversation_id TEXT NOT NULL,
100
+ title TEXT NOT NULL,
101
+ content TEXT NOT NULL,
102
+ timestamp DATETIME NOT NULL,
103
+ FOREIGN KEY (conversation_id) REFERENCES conversation_metadata(conversation_id)
104
+ );
105
+
106
+ -- Table for linking notes to keywords
107
+ CREATE TABLE IF NOT EXISTS rag_qa_note_keywords (
108
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
109
+ note_id INTEGER NOT NULL,
110
+ keyword_id INTEGER NOT NULL,
111
+ FOREIGN KEY (note_id) REFERENCES rag_qa_notes(id),
112
+ FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id)
113
+ );
114
+
115
+ -- Indexes for improved query performance
116
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_chats_conversation_id ON rag_qa_chats(conversation_id);
117
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_chats_timestamp ON rag_qa_chats(timestamp);
118
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_keywords_keyword ON rag_qa_keywords(keyword);
119
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_conversation_keywords_conversation_id ON rag_qa_conversation_keywords(conversation_id);
120
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_conversation_keywords_keyword_id ON rag_qa_conversation_keywords(keyword_id);
121
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_keyword_collections_parent_id ON rag_qa_keyword_collections(parent_id);
122
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_collection_keywords_collection_id ON rag_qa_collection_keywords(collection_id);
123
+ CREATE INDEX IF NOT EXISTS idx_rag_qa_collection_keywords_keyword_id ON rag_qa_collection_keywords(keyword_id);
124
+
125
+ -- Full-text search virtual table for chat content
126
+ CREATE VIRTUAL TABLE IF NOT EXISTS rag_qa_chats_fts USING fts5(conversation_id, timestamp, role, content);
127
+
128
+ -- Trigger to keep the FTS table up to date
129
+ CREATE TRIGGER IF NOT EXISTS rag_qa_chats_ai AFTER INSERT ON rag_qa_chats BEGIN
130
+ INSERT INTO rag_qa_chats_fts(conversation_id, timestamp, role, content) VALUES (new.conversation_id, new.timestamp, new.role, new.content);
131
+ END;
132
+ '''
133
+
134
+ # Database connection management
135
+ @contextmanager
136
+ def get_db_connection():
137
+ conn = sqlite3.connect(rag_qa_db_path)
138
+ try:
139
+ yield conn
140
+ finally:
141
+ conn.close()
142
+
143
+ @contextmanager
144
+ def transaction():
145
+ with get_db_connection() as conn:
146
+ try:
147
+ yield conn
148
+ conn.commit()
149
+ except Exception:
150
+ conn.rollback()
151
+ raise
152
+
153
+ def execute_query(query, params=None, conn=None):
154
+ if conn:
155
+ cursor = conn.cursor()
156
+ if params:
157
+ cursor.execute(query, params)
158
+ else:
159
+ cursor.execute(query)
160
+ return cursor.fetchall()
161
+ else:
162
+ with get_db_connection() as conn:
163
+ cursor = conn.cursor()
164
+ if params:
165
+ cursor.execute(query, params)
166
+ else:
167
+ cursor.execute(query)
168
+ conn.commit()
169
+ return cursor.fetchall()
170
+
171
+ def create_tables():
172
+ with get_db_connection() as conn:
173
+ conn.executescript(SCHEMA_SQL)
174
+ logger.info("All RAG QA Chat tables created successfully")
175
+
176
+ # Initialize the database
177
+ create_tables()
178
+
179
+ #
180
+ # End of Setup
181
+ ############################################################
182
+
183
+
184
+ ############################################################
185
+ #
186
+ # Keyword-related functions
187
+
188
+ # Input validation
189
+ def validate_keyword(keyword):
190
+ if not isinstance(keyword, str):
191
+ raise ValueError("Keyword must be a string")
192
+ if not keyword.strip():
193
+ raise ValueError("Keyword cannot be empty or just whitespace")
194
+ if len(keyword) > 100:
195
+ raise ValueError("Keyword is too long (max 100 characters)")
196
+ if not re.match(r'^[a-zA-Z0-9\s\-_]+$', keyword):
197
+ raise ValueError("Keyword contains invalid characters")
198
+ return keyword.strip()
199
+
200
+ def validate_collection_name(name):
201
+ if not isinstance(name, str):
202
+ raise ValueError("Collection name must be a string")
203
+ if not name.strip():
204
+ raise ValueError("Collection name cannot be empty or just whitespace")
205
+ if len(name) > 100:
206
+ raise ValueError("Collection name is too long (max 100 characters)")
207
+ if not re.match(r'^[a-zA-Z0-9\s\-_]+$', name):
208
+ raise ValueError("Collection name contains invalid characters")
209
+ return name.strip()
210
+
211
+ # Core functions
212
+ def add_keyword(keyword, conn=None):
213
+ try:
214
+ validated_keyword = validate_keyword(keyword)
215
+ query = "INSERT OR IGNORE INTO rag_qa_keywords (keyword) VALUES (?)"
216
+ execute_query(query, (validated_keyword,), conn)
217
+ logger.info(f"Keyword '{validated_keyword}' added successfully")
218
+ except ValueError as e:
219
+ logger.error(f"Invalid keyword: {e}")
220
+ raise
221
+ except Exception as e:
222
+ logger.error(f"Error adding keyword '{keyword}': {e}")
223
+ raise
224
+
225
+ def create_keyword_collection(name, parent_id=None):
226
+ try:
227
+ validated_name = validate_collection_name(name)
228
+ query = "INSERT INTO rag_qa_keyword_collections (name, parent_id) VALUES (?, ?)"
229
+ execute_query(query, (validated_name, parent_id))
230
+ logger.info(f"Keyword collection '{validated_name}' created successfully")
231
+ except ValueError as e:
232
+ logger.error(f"Invalid collection name: {e}")
233
+ raise
234
+ except Exception as e:
235
+ logger.error(f"Error creating keyword collection '{name}': {e}")
236
+ raise
237
+
238
+ def add_keyword_to_collection(collection_name, keyword):
239
+ try:
240
+ validated_collection_name = validate_collection_name(collection_name)
241
+ validated_keyword = validate_keyword(keyword)
242
+
243
+ with transaction() as conn:
244
+ add_keyword(validated_keyword, conn)
245
+
246
+ query = '''
247
+ INSERT INTO rag_qa_collection_keywords (collection_id, keyword_id)
248
+ SELECT c.id, k.id
249
+ FROM rag_qa_keyword_collections c, rag_qa_keywords k
250
+ WHERE c.name = ? AND k.keyword = ?
251
+ '''
252
+ execute_query(query, (validated_collection_name, validated_keyword), conn)
253
+
254
+ logger.info(f"Keyword '{validated_keyword}' added to collection '{validated_collection_name}' successfully")
255
+ except ValueError as e:
256
+ logger.error(f"Invalid input: {e}")
257
+ raise
258
+ except Exception as e:
259
+ logger.error(f"Error adding keyword '{keyword}' to collection '{collection_name}': {e}")
260
+ raise
261
+
262
+ def add_keywords_to_conversation(conversation_id, keywords):
263
+ if not isinstance(keywords, (list, tuple)):
264
+ raise ValueError("Keywords must be a list or tuple")
265
+ try:
266
+ with transaction() as conn:
267
+ for keyword in keywords:
268
+ validated_keyword = validate_keyword(keyword)
269
+ add_keyword(validated_keyword, conn)
270
+
271
+ query = '''
272
+ INSERT INTO rag_qa_conversation_keywords (conversation_id, keyword_id)
273
+ SELECT ?, id FROM rag_qa_keywords WHERE keyword = ?
274
+ '''
275
+ execute_query(query, (conversation_id, validated_keyword), conn)
276
+
277
+ logger.info(f"Keywords added to conversation '{conversation_id}' successfully")
278
+ except ValueError as e:
279
+ logger.error(f"Invalid keyword: {e}")
280
+ raise
281
+ except Exception as e:
282
+ logger.error(f"Error adding keywords to conversation '{conversation_id}': {e}")
283
+ raise
284
+
285
+ def get_keywords_for_conversation(conversation_id):
286
+ try:
287
+ query = '''
288
+ SELECT k.keyword
289
+ FROM rag_qa_keywords k
290
+ JOIN rag_qa_conversation_keywords ck ON k.id = ck.keyword_id
291
+ WHERE ck.conversation_id = ?
292
+ '''
293
+ result = execute_query(query, (conversation_id,))
294
+ keywords = [row[0] for row in result]
295
+ logger.info(f"Retrieved {len(keywords)} keywords for conversation '{conversation_id}'")
296
+ return keywords
297
+ except Exception as e:
298
+ logger.error(f"Error getting keywords for conversation '{conversation_id}': {e}")
299
+ raise
300
+
301
+ def get_keywords_for_collection(collection_name):
302
+ try:
303
+ query = '''
304
+ SELECT k.keyword
305
+ FROM rag_qa_keywords k
306
+ JOIN rag_qa_collection_keywords ck ON k.id = ck.keyword_id
307
+ JOIN rag_qa_keyword_collections c ON ck.collection_id = c.id
308
+ WHERE c.name = ?
309
+ '''
310
+ result = execute_query(query, (collection_name,))
311
+ keywords = [row[0] for row in result]
312
+ logger.info(f"Retrieved {len(keywords)} keywords for collection '{collection_name}'")
313
+ return keywords
314
+ except Exception as e:
315
+ logger.error(f"Error getting keywords for collection '{collection_name}': {e}")
316
+ raise
317
+
318
+ #
319
+ # End of Keyword-related functions
320
+ ###################################################
321
+
322
+
323
+ ###################################################
324
+ #
325
+ # Notes and chat-related functions
326
+
327
+ def save_notes(conversation_id, title, content):
328
+ """Save notes to the database."""
329
+ try:
330
+ query = "INSERT INTO rag_qa_notes (conversation_id, title, content, timestamp) VALUES (?, ?, ?, ?)"
331
+ timestamp = datetime.now().isoformat()
332
+ with transaction() as conn:
333
+ cursor = conn.cursor()
334
+ cursor.execute(query, (conversation_id, title, content, timestamp))
335
+ note_id = cursor.lastrowid
336
+ logger.info(f"Notes saved for conversation '{conversation_id}', note ID '{note_id}'")
337
+ return note_id
338
+ except Exception as e:
339
+ logger.error(f"Error saving notes for conversation '{conversation_id}': {e}")
340
+ raise
341
+
342
+ def update_note(note_id, title, content):
343
+ try:
344
+ query = "UPDATE rag_qa_notes SET title = ?, content = ?, timestamp = ? WHERE id = ?"
345
+ timestamp = datetime.now().isoformat()
346
+ execute_query(query, (title, content, timestamp, note_id))
347
+ logger.info(f"Note ID '{note_id}' updated successfully")
348
+ except Exception as e:
349
+ logger.error(f"Error updating note ID '{note_id}': {e}")
350
+ raise
351
+
352
+ def get_notes(conversation_id):
353
+ """Retrieve notes for a given conversation."""
354
+ try:
355
+ query = "SELECT content FROM rag_qa_notes WHERE conversation_id = ?"
356
+ result = execute_query(query, (conversation_id,))
357
+ notes = [row[0] for row in result]
358
+ logger.info(f"Retrieved {len(notes)} notes for conversation '{conversation_id}'")
359
+ return notes
360
+ except Exception as e:
361
+ logger.error(f"Error getting notes for conversation '{conversation_id}': {e}")
362
+ raise
363
+
364
+ def get_note_by_id(note_id):
365
+ try:
366
+ query = "SELECT id, title, content FROM rag_qa_notes WHERE id = ?"
367
+ result = execute_query(query, (note_id,))
368
+ return result
369
+ except Exception as e:
370
+ logger.error(f"Error getting note by ID '{note_id}': {e}")
371
+ raise
372
+
373
+ def get_notes_by_keywords(keywords, page=1, page_size=20):
374
+ try:
375
+ placeholders = ','.join(['?'] * len(keywords))
376
+ query = f'''
377
+ SELECT n.id, n.title, n.content, n.timestamp
378
+ FROM rag_qa_notes n
379
+ JOIN rag_qa_note_keywords nk ON n.id = nk.note_id
380
+ JOIN rag_qa_keywords k ON nk.keyword_id = k.id
381
+ WHERE k.keyword IN ({placeholders})
382
+ ORDER BY n.timestamp DESC
383
+ '''
384
+ results, total_pages, total_count = get_paginated_results(query, tuple(keywords), page, page_size)
385
+ logger.info(f"Retrieved {len(results)} notes matching keywords: {', '.join(keywords)} (page {page} of {total_pages})")
386
+ notes = [(row[0], row[1], row[2], row[3]) for row in results]
387
+ return notes, total_pages, total_count
388
+ except Exception as e:
389
+ logger.error(f"Error getting notes by keywords: {e}")
390
+ raise
391
+
392
+ def get_notes_by_keyword_collection(collection_name, page=1, page_size=20):
393
+ try:
394
+ query = '''
395
+ SELECT n.id, n.title, n.content, n.timestamp
396
+ FROM rag_qa_notes n
397
+ JOIN rag_qa_note_keywords nk ON n.id = nk.note_id
398
+ JOIN rag_qa_keywords k ON nk.keyword_id = k.id
399
+ JOIN rag_qa_collection_keywords ck ON k.id = ck.keyword_id
400
+ JOIN rag_qa_keyword_collections c ON ck.collection_id = c.id
401
+ WHERE c.name = ?
402
+ ORDER BY n.timestamp DESC
403
+ '''
404
+ results, total_pages, total_count = get_paginated_results(query, (collection_name,), page, page_size)
405
+ logger.info(f"Retrieved {len(results)} notes for collection '{collection_name}' (page {page} of {total_pages})")
406
+ notes = [(row[0], row[1], row[2], row[3]) for row in results]
407
+ return notes, total_pages, total_count
408
+ except Exception as e:
409
+ logger.error(f"Error getting notes by keyword collection '{collection_name}': {e}")
410
+ raise
411
+
412
+ def clear_notes(conversation_id):
413
+ """Clear all notes for a given conversation."""
414
+ try:
415
+ query = "DELETE FROM rag_qa_notes WHERE conversation_id = ?"
416
+ execute_query(query, (conversation_id,))
417
+ logger.info(f"Cleared notes for conversation '{conversation_id}'")
418
+ except Exception as e:
419
+ logger.error(f"Error clearing notes for conversation '{conversation_id}': {e}")
420
+ raise
421
+
422
+ def add_keywords_to_note(note_id, keywords):
423
+ """Associate keywords with a note."""
424
+ try:
425
+ with transaction() as conn:
426
+ for keyword in keywords:
427
+ validated_keyword = validate_keyword(keyword)
428
+ add_keyword(validated_keyword, conn)
429
+
430
+ # Retrieve the keyword ID
431
+ query = "SELECT id FROM rag_qa_keywords WHERE keyword = ?"
432
+ result = execute_query(query, (validated_keyword,), conn)
433
+ if result:
434
+ keyword_id = result[0][0]
435
+ else:
436
+ raise Exception(f"Keyword '{validated_keyword}' not found after insertion")
437
+
438
+ # Link the note and keyword
439
+ query = "INSERT INTO rag_qa_note_keywords (note_id, keyword_id) VALUES (?, ?)"
440
+ execute_query(query, (note_id, keyword_id), conn)
441
+
442
+ logger.info(f"Keywords added to note ID '{note_id}' successfully")
443
+ except Exception as e:
444
+ logger.error(f"Error adding keywords to note ID '{note_id}': {e}")
445
+ raise
446
+
447
+ def get_keywords_for_note(note_id):
448
+ """Retrieve keywords associated with a given note."""
449
+ try:
450
+ query = '''
451
+ SELECT k.keyword
452
+ FROM rag_qa_keywords k
453
+ JOIN rag_qa_note_keywords nk ON k.id = nk.keyword_id
454
+ WHERE nk.note_id = ?
455
+ '''
456
+ result = execute_query(query, (note_id,))
457
+ keywords = [row[0] for row in result]
458
+ logger.info(f"Retrieved {len(keywords)} keywords for note ID '{note_id}'")
459
+ return keywords
460
+ except Exception as e:
461
+ logger.error(f"Error getting keywords for note ID '{note_id}': {e}")
462
+ raise
463
+
464
+ def clear_keywords_from_note(note_id):
465
+ """Clear all keywords from a given note."""
466
+ try:
467
+ query = "DELETE FROM rag_qa_note_keywords WHERE note_id = ?"
468
+ execute_query(query, (note_id,))
469
+ logger.info(f"Cleared keywords for note ID '{note_id}'")
470
+ except Exception as e:
471
+ logger.error(f"Error clearing keywords for note ID '{note_id}': {e}")
472
+ raise
473
+
474
+ def delete_note_by_id(note_id, conn=None):
475
+ """Delete a note and its associated keywords."""
476
+ try:
477
+ # Delete note keywords
478
+ execute_query("DELETE FROM rag_qa_note_keywords WHERE note_id = ?", (note_id,), conn)
479
+ # Delete the note
480
+ execute_query("DELETE FROM rag_qa_notes WHERE id = ?", (note_id,), conn)
481
+ logging.info(f"Note ID '{note_id}' deleted successfully.")
482
+ except Exception as e:
483
+ logger.error(f"Error deleting note ID '{note_id}': {e}")
484
+ raise
485
+
486
+ def delete_note(note_id):
487
+ """Delete a note by ID."""
488
+ try:
489
+ with transaction() as conn:
490
+ delete_note_by_id(note_id, conn)
491
+ except Exception as e:
492
+ logger.error(f"Error deleting note ID '{note_id}': {e}")
493
+ raise
494
+
495
+ #
496
+ # End of Notes related functions
497
+ ###################################################
498
+
499
+
500
+ ###################################################
501
+ #
502
+ # Chat-related functions
503
+
504
+ def save_message(conversation_id, role, content):
505
+ try:
506
+ timestamp = datetime.now().isoformat()
507
+ query = "INSERT INTO rag_qa_chats (conversation_id, timestamp, role, content) VALUES (?, ?, ?, ?)"
508
+ execute_query(query, (conversation_id, timestamp, role, content))
509
+
510
+ # Update last_updated in conversation_metadata
511
+ update_query = "UPDATE conversation_metadata SET last_updated = ? WHERE conversation_id = ?"
512
+ execute_query(update_query, (timestamp, conversation_id))
513
+
514
+ logger.info(f"Message saved for conversation '{conversation_id}'")
515
+ except Exception as e:
516
+ logger.error(f"Error saving message for conversation '{conversation_id}': {e}")
517
+ raise
518
+
519
+ def start_new_conversation(title="Untitled Conversation"):
520
+ try:
521
+ conversation_id = str(uuid.uuid4())
522
+ query = "INSERT INTO conversation_metadata (conversation_id, created_at, last_updated, title) VALUES (?, ?, ?, ?)"
523
+ now = datetime.now().isoformat()
524
+ execute_query(query, (conversation_id, now, now, title))
525
+ logger.info(f"New conversation '{conversation_id}' started with title '{title}'")
526
+ return conversation_id
527
+ except Exception as e:
528
+ logger.error(f"Error starting new conversation: {e}")
529
+ raise
530
+
531
+ def get_all_conversations(page=1, page_size=20):
532
+ try:
533
+ query = "SELECT conversation_id, title FROM conversation_metadata ORDER BY last_updated DESC"
534
+ results, total_pages, total_count = get_paginated_results(query, page=page, page_size=page_size)
535
+ conversations = [(row[0], row[1]) for row in results]
536
+ logger.info(f"Retrieved {len(conversations)} conversations (page {page} of {total_pages})")
537
+ return conversations, total_pages, total_count
538
+ except Exception as e:
539
+ logger.error(f"Error getting conversations: {e}")
540
+ raise
541
+
542
+ # Pagination helper function
543
+ def get_paginated_results(query, params=None, page=1, page_size=20):
544
+ try:
545
+ offset = (page - 1) * page_size
546
+ paginated_query = f"{query} LIMIT ? OFFSET ?"
547
+ if params:
548
+ paginated_params = params + (page_size, offset)
549
+ else:
550
+ paginated_params = (page_size, offset)
551
+
552
+ result = execute_query(paginated_query, paginated_params)
553
+
554
+ count_query = f"SELECT COUNT(*) FROM ({query}) AS total"
555
+ count_params = params if params else ()
556
+
557
+ total_count = execute_query(count_query, count_params)[0][0]
558
+
559
+ total_pages = (total_count + page_size - 1) // page_size
560
+
561
+ logger.info(f"Retrieved page {page} of {total_pages} (total items: {total_count})")
562
+ return result, total_pages, total_count
563
+ except Exception as e:
564
+ logger.error(f"Error retrieving paginated results: {e}")
565
+ raise
566
+
567
+ def get_all_collections(page=1, page_size=20):
568
+ try:
569
+ query = "SELECT name FROM rag_qa_keyword_collections"
570
+ results, total_pages, total_count = get_paginated_results(query, page=page, page_size=page_size)
571
+ collections = [row[0] for row in results]
572
+ logger.info(f"Retrieved {len(collections)} keyword collections (page {page} of {total_pages})")
573
+ return collections, total_pages, total_count
574
+ except Exception as e:
575
+ logger.error(f"Error getting collections: {e}")
576
+ raise
577
+
578
+ def search_conversations_by_keywords(keywords, page=1, page_size=20):
579
+ try:
580
+ placeholders = ','.join(['?' for _ in keywords])
581
+ query = f'''
582
+ SELECT DISTINCT cm.conversation_id, cm.title
583
+ FROM conversation_metadata cm
584
+ JOIN rag_qa_conversation_keywords ck ON cm.conversation_id = ck.conversation_id
585
+ JOIN rag_qa_keywords k ON ck.keyword_id = k.id
586
+ WHERE k.keyword IN ({placeholders})
587
+ '''
588
+ results, total_pages, total_count = get_paginated_results(query, tuple(keywords), page, page_size)
589
+ logger.info(
590
+ f"Found {total_count} conversations matching keywords: {', '.join(keywords)} (page {page} of {total_pages})")
591
+ return results, total_pages, total_count
592
+ except Exception as e:
593
+ logger.error(f"Error searching conversations by keywords {keywords}: {e}")
594
+ raise
595
+
596
+ def load_chat_history(conversation_id, page=1, page_size=50):
597
+ try:
598
+ query = "SELECT role, content FROM rag_qa_chats WHERE conversation_id = ? ORDER BY timestamp"
599
+ results, total_pages, total_count = get_paginated_results(query, (conversation_id,), page, page_size)
600
+ logger.info(
601
+ f"Loaded {len(results)} messages for conversation '{conversation_id}' (page {page} of {total_pages})")
602
+ return results, total_pages, total_count
603
+ except Exception as e:
604
+ logger.error(f"Error loading chat history for conversation '{conversation_id}': {e}")
605
+ raise
606
+
607
+ def update_conversation_title(conversation_id, new_title):
608
+ """Update the title of a conversation."""
609
+ try:
610
+ query = "UPDATE conversation_metadata SET title = ? WHERE conversation_id = ?"
611
+ execute_query(query, (new_title, conversation_id))
612
+ logger.info(f"Conversation '{conversation_id}' title updated to '{new_title}'")
613
+ except Exception as e:
614
+ logger.error(f"Error updating conversation title: {e}")
615
+ raise
616
+
617
+ def delete_conversation(conversation_id):
618
+ """Delete a conversation and its associated messages and notes."""
619
+ try:
620
+ with transaction() as conn:
621
+ # Delete messages
622
+ execute_query("DELETE FROM rag_qa_chats WHERE conversation_id = ?", (conversation_id,), conn)
623
+ # Delete conversation metadata
624
+ execute_query("DELETE FROM conversation_metadata WHERE conversation_id = ?", (conversation_id,), conn)
625
+ # Delete conversation keywords
626
+ execute_query("DELETE FROM rag_qa_conversation_keywords WHERE conversation_id = ?", (conversation_id,), conn)
627
+ # Delete notes associated with the conversation
628
+ note_ids = execute_query("SELECT id FROM rag_qa_notes WHERE conversation_id = ?", (conversation_id,), conn)
629
+ for (note_id,) in note_ids:
630
+ delete_note_by_id(note_id, conn)
631
+ logging.info(f"Conversation '{conversation_id}' deleted successfully.")
632
+ except Exception as e:
633
+ logger.error(f"Error deleting conversation '{conversation_id}': {e}")
634
+ raise
635
+
636
+ #
637
+ # End of Chat-related functions
638
+ ###################################################
639
+
640
+
641
+ ###################################################
642
+ #
643
+ # Functions to export DB data
644
+
645
+ def fetch_all_conversations():
646
+ try:
647
+ # Fetch all conversation IDs and titles
648
+ query = "SELECT conversation_id, title FROM conversation_metadata ORDER BY last_updated DESC"
649
+ results = execute_query(query)
650
+ conversations = []
651
+ for row in results:
652
+ conversation_id, title = row
653
+ # Fetch all messages for this conversation
654
+ messages = load_all_chat_history(conversation_id)
655
+ conversations.append((conversation_id, title, messages))
656
+ logger.info(f"Fetched all conversations: {len(conversations)} found.")
657
+ return conversations
658
+ except Exception as e:
659
+ logger.error(f"Error fetching all conversations: {e}")
660
+ raise
661
+
662
+ def load_all_chat_history(conversation_id):
663
+ try:
664
+ query = "SELECT role, content FROM rag_qa_chats WHERE conversation_id = ? ORDER BY timestamp"
665
+ results = execute_query(query, (conversation_id,))
666
+ messages = [(row[0], row[1]) for row in results]
667
+ return messages
668
+ except Exception as e:
669
+ logger.error(f"Error loading chat history for conversation '{conversation_id}': {e}")
670
+ raise
671
+
672
+ def fetch_all_notes():
673
+ try:
674
+ query = "SELECT id, title, content FROM rag_qa_notes ORDER BY timestamp DESC"
675
+ results = execute_query(query)
676
+ notes = [(row[0], row[1], row[2]) for row in results]
677
+ logger.info(f"Fetched all notes: {len(notes)} found.")
678
+ return notes
679
+ except Exception as e:
680
+ logger.error(f"Error fetching all notes: {e}")
681
+ raise
682
+
683
+ def fetch_conversations_by_ids(conversation_ids):
684
+ try:
685
+ if not conversation_ids:
686
+ return []
687
+ placeholders = ','.join(['?'] * len(conversation_ids))
688
+ query = f"SELECT conversation_id, title FROM conversation_metadata WHERE conversation_id IN ({placeholders})"
689
+ results = execute_query(query, conversation_ids)
690
+ conversations = []
691
+ for row in results:
692
+ conversation_id, title = row
693
+ # Fetch all messages for this conversation
694
+ messages = load_all_chat_history(conversation_id)
695
+ conversations.append((conversation_id, title, messages))
696
+ logger.info(f"Fetched {len(conversations)} conversations by IDs.")
697
+ return conversations
698
+ except Exception as e:
699
+ logger.error(f"Error fetching conversations by IDs: {e}")
700
+ raise
701
+
702
+ def fetch_notes_by_ids(note_ids):
703
+ try:
704
+ if not note_ids:
705
+ return []
706
+ placeholders = ','.join(['?'] * len(note_ids))
707
+ query = f"SELECT id, title, content FROM rag_qa_notes WHERE id IN ({placeholders})"
708
+ results = execute_query(query, note_ids)
709
+ notes = [(row[0], row[1], row[2]) for row in results]
710
+ logger.info(f"Fetched {len(notes)} notes by IDs.")
711
+ return notes
712
+ except Exception as e:
713
+ logger.error(f"Error fetching notes by IDs: {e}")
714
+ raise
715
+
716
+ #
717
+ # End of Export functions
718
+ ###################################################
719
+
720
+ #
721
+ # End of RAG_QA_Chat_DB.py
722
+ ####################################################################################################
App_Function_Libraries/DB/SQLite_DB.py ADDED
The diff for this file is too large to render. See raw diff
 
App_Function_Libraries/DB/__init__.py ADDED
File without changes
App_Function_Libraries/Gradio_Related.py ADDED
@@ -0,0 +1,420 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Gradio_Related.py
2
+ #########################################
3
+ # Gradio UI Functions Library
4
+ # I fucking hate Gradio.
5
+ #
6
+ #########################################
7
+ #
8
+ # Built-In Imports
9
+ import logging
10
+ import os
11
+ import webbrowser
12
+
13
+ #
14
+ # Import 3rd-Party Libraries
15
+ import gradio as gr
16
+ #
17
+ # Local Imports
18
+ from App_Function_Libraries.DB.DB_Manager import get_db_config
19
+ from App_Function_Libraries.Gradio_UI.Arxiv_tab import create_arxiv_tab
20
+ from App_Function_Libraries.Gradio_UI.Audio_ingestion_tab import create_audio_processing_tab
21
+ from App_Function_Libraries.Gradio_UI.Book_Ingestion_tab import create_import_book_tab
22
+ from App_Function_Libraries.Gradio_UI.Character_Chat_tab import create_character_card_interaction_tab, create_character_chat_mgmt_tab, create_custom_character_card_tab, \
23
+ create_character_card_validation_tab, create_export_characters_tab
24
+ from App_Function_Libraries.Gradio_UI.Character_interaction_tab import create_narrator_controlled_conversation_tab, \
25
+ create_multiple_character_chat_tab
26
+ from App_Function_Libraries.Gradio_UI.Chat_ui import create_chat_management_tab, \
27
+ create_chat_interface_four, create_chat_interface_multi_api, create_chat_interface_stacked, create_chat_interface
28
+ from App_Function_Libraries.Gradio_UI.Config_tab import create_config_editor_tab
29
+ from App_Function_Libraries.Gradio_UI.Explain_summarize_tab import create_summarize_explain_tab
30
+ from App_Function_Libraries.Gradio_UI.Export_Functionality import create_export_tab
31
+ from App_Function_Libraries.Gradio_UI.Backup_Functionality import create_backup_tab, create_view_backups_tab, \
32
+ create_restore_backup_tab
33
+ from App_Function_Libraries.Gradio_UI.Import_Functionality import create_import_single_prompt_tab, \
34
+ create_import_obsidian_vault_tab, create_import_item_tab, create_import_multiple_prompts_tab
35
+ from App_Function_Libraries.Gradio_UI.Introduction_tab import create_introduction_tab
36
+ from App_Function_Libraries.Gradio_UI.Keywords import create_view_keywords_tab, create_add_keyword_tab, \
37
+ create_delete_keyword_tab, create_export_keywords_tab
38
+ from App_Function_Libraries.Gradio_UI.Live_Recording import create_live_recording_tab
39
+ from App_Function_Libraries.Gradio_UI.Llamafile_tab import create_chat_with_llamafile_tab
40
+ #from App_Function_Libraries.Gradio_UI.MMLU_Pro_tab import create_mmlu_pro_tab
41
+ from App_Function_Libraries.Gradio_UI.Media_edit import create_prompt_clone_tab, create_prompt_edit_tab, \
42
+ create_media_edit_and_clone_tab, create_media_edit_tab
43
+ from App_Function_Libraries.Gradio_UI.Media_wiki_tab import create_mediawiki_import_tab, create_mediawiki_config_tab
44
+ from App_Function_Libraries.Gradio_UI.PDF_ingestion_tab import create_pdf_ingestion_tab, create_pdf_ingestion_test_tab
45
+ from App_Function_Libraries.Gradio_UI.Plaintext_tab_import import create_plain_text_import_tab
46
+ from App_Function_Libraries.Gradio_UI.Podcast_tab import create_podcast_tab
47
+ from App_Function_Libraries.Gradio_UI.Prompt_Suggestion_tab import create_prompt_suggestion_tab
48
+ from App_Function_Libraries.Gradio_UI.RAG_QA_Chat_tab import create_rag_qa_chat_tab, create_rag_qa_notes_management_tab, \
49
+ create_rag_qa_chat_management_tab
50
+ from App_Function_Libraries.Gradio_UI.Re_summarize_tab import create_resummary_tab
51
+ from App_Function_Libraries.Gradio_UI.Search_Tab import create_prompt_search_tab, \
52
+ create_search_summaries_tab, create_search_tab
53
+ from App_Function_Libraries.Gradio_UI.RAG_Chat_tab import create_rag_tab
54
+ from App_Function_Libraries.Gradio_UI.Embeddings_tab import create_embeddings_tab, create_view_embeddings_tab, \
55
+ create_purge_embeddings_tab
56
+ from App_Function_Libraries.Gradio_UI.Trash import create_view_trash_tab, create_empty_trash_tab, \
57
+ create_delete_trash_tab, create_search_and_mark_trash_tab
58
+ from App_Function_Libraries.Gradio_UI.Utilities import create_utilities_yt_timestamp_tab, create_utilities_yt_audio_tab, \
59
+ create_utilities_yt_video_tab
60
+ from App_Function_Libraries.Gradio_UI.Video_transcription_tab import create_video_transcription_tab
61
+ from App_Function_Libraries.Gradio_UI.View_tab import create_manage_items_tab
62
+ from App_Function_Libraries.Gradio_UI.Website_scraping_tab import create_website_scraping_tab
63
+ from App_Function_Libraries.Gradio_UI.Chat_Workflows import chat_workflows_tab
64
+ from App_Function_Libraries.Gradio_UI.View_DB_Items_tab import create_prompt_view_tab, \
65
+ create_view_all_with_versions_tab, create_viewing_tab
66
+ #
67
+ # Gradio UI Imports
68
+ from App_Function_Libraries.Gradio_UI.Evaluations_Benchmarks_tab import create_geval_tab, create_infinite_bench_tab
69
+ #from App_Function_Libraries.Local_LLM.Local_LLM_huggingface import create_huggingface_tab
70
+ from App_Function_Libraries.Local_LLM.Local_LLM_ollama import create_ollama_tab
71
+ #
72
+ #######################################################################################################################
73
+ # Function Definitions
74
+ #
75
+
76
+
77
+ # Disable Gradio Analytics
78
+ os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
79
+
80
+
81
+ custom_prompt_input = None
82
+ server_mode = False
83
+ share_public = False
84
+ custom_prompt_summarize_bulleted_notes = ("""
85
+ <s>You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST]
86
+ **Bulleted Note Creation Guidelines**
87
+
88
+ **Headings**:
89
+ - Based on referenced topics, not categories like quotes or terms
90
+ - Surrounded by **bold** formatting
91
+ - Not listed as bullet points
92
+ - No space between headings and list items underneath
93
+
94
+ **Emphasis**:
95
+ - **Important terms** set in bold font
96
+ - **Text ending in a colon**: also bolded
97
+
98
+ **Review**:
99
+ - Ensure adherence to specified format
100
+ - Do not reference these instructions in your response.</s>[INST] {{ .Prompt }} [/INST]
101
+ """)
102
+ #
103
+ # End of globals
104
+ #######################################################################################################################
105
+ #
106
+ # Start of Video/Audio Transcription and Summarization Functions
107
+ #
108
+ # Functions:
109
+ # FIXME
110
+ #
111
+ #
112
+ ################################################################################################################
113
+ # Functions for Re-Summarization
114
+ #
115
+ # Functions:
116
+ # FIXME
117
+ # End of Re-Summarization Functions
118
+ #
119
+ ############################################################################################################################################################################################################################
120
+ #
121
+ # Explain/Summarize This Tab
122
+ #
123
+ # Functions:
124
+ # FIXME
125
+ #
126
+ #
127
+ ############################################################################################################################################################################################################################
128
+ #
129
+ # Transcript Comparison Tab
130
+ #
131
+ # Functions:
132
+ # FIXME
133
+ #
134
+ #
135
+ ###########################################################################################################################################################################################################################
136
+ #
137
+ # Search Tab
138
+ #
139
+ # Functions:
140
+ # FIXME
141
+ #
142
+ # End of Search Tab Functions
143
+ #
144
+ ##############################################################################################################################################################################################################################
145
+ #
146
+ # Llamafile Tab
147
+ #
148
+ # Functions:
149
+ # FIXME
150
+ #
151
+ # End of Llamafile Tab Functions
152
+ ##############################################################################################################################################################################################################################
153
+ #
154
+ # Chat Interface Tab Functions
155
+ #
156
+ # Functions:
157
+ # FIXME
158
+ #
159
+ #
160
+ # End of Chat Interface Tab Functions
161
+ ################################################################################################################################################################################################################################
162
+ #
163
+ # Media Edit Tab Functions
164
+ # Functions:
165
+ # Fixme
166
+ # create_media_edit_tab():
167
+ ##### Trash Tab
168
+ # FIXME
169
+ # Functions:
170
+ #
171
+ # End of Media Edit Tab Functions
172
+ ################################################################################################################
173
+ #
174
+ # Import Items Tab Functions
175
+ #
176
+ # Functions:
177
+ #FIXME
178
+ # End of Import Items Tab Functions
179
+ ################################################################################################################
180
+ #
181
+ # Export Items Tab Functions
182
+ #
183
+ # Functions:
184
+ # FIXME
185
+ #
186
+ #
187
+ # End of Export Items Tab Functions
188
+ ################################################################################################################
189
+ #
190
+ # Keyword Management Tab Functions
191
+ #
192
+ # Functions:
193
+ # create_view_keywords_tab():
194
+ # FIXME
195
+ #
196
+ # End of Keyword Management Tab Functions
197
+ ################################################################################################################
198
+ #
199
+ # Document Editing Tab Functions
200
+ #
201
+ # Functions:
202
+ # #FIXME
203
+ #
204
+ #
205
+ ################################################################################################################
206
+ #
207
+ # Utilities Tab Functions
208
+ # Functions:
209
+ # create_utilities_yt_video_tab():
210
+ # #FIXME
211
+
212
+ #
213
+ # End of Utilities Tab Functions
214
+ ################################################################################################################
215
+
216
+ # FIXME - Prompt sample box
217
+ #
218
+ # # Sample data
219
+ # prompts_category_1 = [
220
+ # "What are the key points discussed in the video?",
221
+ # "Summarize the main arguments made by the speaker.",
222
+ # "Describe the conclusions of the study presented."
223
+ # ]
224
+ #
225
+ # prompts_category_2 = [
226
+ # "How does the proposed solution address the problem?",
227
+ # "What are the implications of the findings?",
228
+ # "Can you explain the theory behind the observed phenomenon?"
229
+ # ]
230
+ #
231
+ # all_prompts2 = prompts_category_1 + prompts_category_2
232
+
233
+
234
+ def launch_ui(share_public=None, server_mode=False):
235
+ webbrowser.open_new_tab('http://127.0.0.1:7860/?__theme=dark')
236
+ share=share_public
237
+ css = """
238
+ .result-box {
239
+ margin-bottom: 20px;
240
+ border: 1px solid #ddd;
241
+ padding: 10px;
242
+ }
243
+ .result-box.error {
244
+ border-color: #ff0000;
245
+ background-color: #ffeeee;
246
+ }
247
+ .transcription, .summary {
248
+ max-height: 800px;
249
+ overflow-y: auto;
250
+ border: 1px solid #eee;
251
+ padding: 10px;
252
+ margin-top: 10px;
253
+ }
254
+ """
255
+
256
+ with gr.Blocks(theme='bethecloud/storj_theme',css=css) as iface:
257
+ gr.HTML(
258
+ """
259
+ <script>
260
+ document.addEventListener('DOMContentLoaded', (event) => {
261
+ document.body.classList.add('dark');
262
+ document.querySelector('gradio-app').style.backgroundColor = 'var(--color-background-primary)';
263
+ });
264
+ </script>
265
+ """
266
+ )
267
+ db_config = get_db_config()
268
+ db_type = db_config['type']
269
+ gr.Markdown(f"# tl/dw: Your LLM-powered Research Multi-tool")
270
+ gr.Markdown(f"(Using {db_type.capitalize()} Database)")
271
+ with gr.Tabs():
272
+ with gr.TabItem("Transcription / Summarization / Ingestion", id="ingestion-grouping", visible=True):
273
+ with gr.Tabs():
274
+ create_video_transcription_tab()
275
+ create_audio_processing_tab()
276
+ create_podcast_tab()
277
+ create_import_book_tab()
278
+ create_plain_text_import_tab()
279
+ create_website_scraping_tab()
280
+ create_pdf_ingestion_tab()
281
+ create_pdf_ingestion_test_tab()
282
+ create_resummary_tab()
283
+ create_summarize_explain_tab()
284
+ create_live_recording_tab()
285
+ create_arxiv_tab()
286
+
287
+ with gr.TabItem("Text Search", id="text search", visible=True):
288
+ create_search_tab()
289
+ create_search_summaries_tab()
290
+
291
+ with gr.TabItem("RAG Chat/Search", id="RAG Chat Notes group", visible=True):
292
+ create_rag_tab()
293
+ create_rag_qa_chat_tab()
294
+ create_rag_qa_notes_management_tab()
295
+ create_rag_qa_chat_management_tab()
296
+
297
+ with gr.TabItem("Chat with an LLM", id="LLM Chat group", visible=True):
298
+ create_chat_interface()
299
+ create_chat_interface_stacked()
300
+ create_chat_interface_multi_api()
301
+ create_chat_interface_four()
302
+ create_chat_with_llamafile_tab()
303
+ create_chat_management_tab()
304
+ chat_workflows_tab()
305
+
306
+
307
+ with gr.TabItem("Character Chat", id="character chat group", visible=True):
308
+ create_character_card_interaction_tab()
309
+ create_character_chat_mgmt_tab()
310
+ create_custom_character_card_tab()
311
+ create_character_card_validation_tab()
312
+ create_multiple_character_chat_tab()
313
+ create_narrator_controlled_conversation_tab()
314
+ create_export_characters_tab()
315
+
316
+
317
+ with gr.TabItem("View DB Items", id="view db items group", visible=True):
318
+ # This one works
319
+ create_view_all_with_versions_tab()
320
+ # This one is WIP
321
+ create_viewing_tab()
322
+ create_prompt_view_tab()
323
+
324
+
325
+ with gr.TabItem("Prompts", id='view prompts group', visible=True):
326
+ create_prompt_view_tab()
327
+ create_prompt_search_tab()
328
+ create_prompt_edit_tab()
329
+ create_prompt_clone_tab()
330
+ create_prompt_suggestion_tab()
331
+
332
+
333
+ with gr.TabItem("Manage / Edit Existing Items", id="manage group", visible=True):
334
+ create_media_edit_tab()
335
+ create_manage_items_tab()
336
+ create_media_edit_and_clone_tab()
337
+ # FIXME
338
+ #create_compare_transcripts_tab()
339
+
340
+
341
+ with gr.TabItem("Embeddings Management", id="embeddings group", visible=True):
342
+ create_embeddings_tab()
343
+ create_view_embeddings_tab()
344
+ create_purge_embeddings_tab()
345
+
346
+ with gr.TabItem("Writing Tools", id="writing_tools group", visible=True):
347
+ from App_Function_Libraries.Gradio_UI.Writing_tab import create_document_feedback_tab
348
+ create_document_feedback_tab()
349
+ from App_Function_Libraries.Gradio_UI.Writing_tab import create_grammar_style_check_tab
350
+ create_grammar_style_check_tab()
351
+ from App_Function_Libraries.Gradio_UI.Writing_tab import create_tone_adjustment_tab
352
+ create_tone_adjustment_tab()
353
+ from App_Function_Libraries.Gradio_UI.Writing_tab import create_creative_writing_tab
354
+ create_creative_writing_tab()
355
+ from App_Function_Libraries.Gradio_UI.Writing_tab import create_mikupad_tab
356
+ create_mikupad_tab()
357
+
358
+
359
+ with gr.TabItem("Keywords", id="keywords group", visible=True):
360
+ create_view_keywords_tab()
361
+ create_add_keyword_tab()
362
+ create_delete_keyword_tab()
363
+ create_export_keywords_tab()
364
+
365
+ with gr.TabItem("Import", id="import group", visible=True):
366
+ create_import_item_tab()
367
+ create_import_obsidian_vault_tab()
368
+ create_import_single_prompt_tab()
369
+ create_import_multiple_prompts_tab()
370
+ create_mediawiki_import_tab()
371
+ create_mediawiki_config_tab()
372
+
373
+ with gr.TabItem("Export", id="export group", visible=True):
374
+ create_export_tab()
375
+
376
+ with gr.TabItem("Backup Management", id="backup group", visible=True):
377
+ create_backup_tab()
378
+ create_view_backups_tab()
379
+ create_restore_backup_tab()
380
+
381
+ with gr.TabItem("Utilities", id="util group", visible=True):
382
+ create_utilities_yt_video_tab()
383
+ create_utilities_yt_audio_tab()
384
+ create_utilities_yt_timestamp_tab()
385
+
386
+ with gr.TabItem("Local LLM", id="local llm group", visible=True):
387
+ create_chat_with_llamafile_tab()
388
+ create_ollama_tab()
389
+ #create_huggingface_tab()
390
+
391
+ with gr.TabItem("Trashcan", id="trashcan group", visible=True):
392
+ create_search_and_mark_trash_tab()
393
+ create_view_trash_tab()
394
+ create_delete_trash_tab()
395
+ create_empty_trash_tab()
396
+
397
+ with gr.TabItem("Evaluations", id="eval", visible=True):
398
+ create_geval_tab()
399
+ create_infinite_bench_tab()
400
+ # FIXME
401
+ #create_mmlu_pro_tab()
402
+
403
+ with gr.TabItem("Introduction/Help", id="introduction group", visible=True):
404
+ create_introduction_tab()
405
+
406
+ with gr.TabItem("Config Editor", id="config group"):
407
+ create_config_editor_tab()
408
+
409
+ # Launch the interface
410
+ server_port_variable = 7860
411
+ os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False'
412
+ if share==True:
413
+ iface.launch(share=True)
414
+ elif server_mode and not share_public:
415
+ iface.launch(share=False, server_name="0.0.0.0", server_port=server_port_variable, )
416
+ else:
417
+ try:
418
+ iface.launch(share=False, server_name="0.0.0.0", server_port=server_port_variable, )
419
+ except Exception as e:
420
+ logging.error(f"Error launching interface: {str(e)}")
App_Function_Libraries/Gradio_UI/Arxiv_tab.py ADDED
@@ -0,0 +1,230 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Arxiv_tab.py
2
+ # Description: This file contains the Gradio UI for searching, browsing, and ingesting arXiv papers.
3
+ #
4
+ # Imports
5
+ import tempfile
6
+ from datetime import datetime
7
+ import requests
8
+
9
+ from App_Function_Libraries.PDF.PDF_Ingestion_Lib import extract_text_and_format_from_pdf
10
+ #
11
+ # Local Imports
12
+ from App_Function_Libraries.Third_Party.Arxiv import convert_xml_to_markdown, fetch_arxiv_xml, parse_arxiv_feed, \
13
+ build_query_url, ARXIV_PAGE_SIZE, fetch_arxiv_pdf_url
14
+ from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords
15
+ #
16
+ import gradio as gr
17
+ #
18
+ #####################################################################################################
19
+ #
20
+ # Functions:
21
+
22
+ def create_arxiv_tab():
23
+ with gr.TabItem("Arxiv Search & Ingest", visible=True):
24
+ gr.Markdown("# arXiv Search, Browse, Download, and Ingest")
25
+ gr.Markdown("#### Thank you to arXiv for use of its open access interoperability.")
26
+ with gr.Row():
27
+ with gr.Column(scale=1):
28
+ # Search Inputs
29
+ with gr.Row():
30
+ with gr.Column():
31
+ search_query = gr.Textbox(label="Search Query", placeholder="e.g., machine learning")
32
+ author_filter = gr.Textbox(label="Author", placeholder="e.g., John Doe")
33
+ year_filter = gr.Number(label="Year", precision=0)
34
+ search_button = gr.Button("Search")
35
+
36
+ with gr.Column(scale=2):
37
+ # Pagination Controls
38
+ paper_selector = gr.Radio(label="Select a Paper", choices=[], interactive=True)
39
+ prev_button = gr.Button("Previous Page")
40
+ next_button = gr.Button("Next Page")
41
+ page_info = gr.Textbox(label="Page", value="1", interactive=False)
42
+
43
+ # Ingestion Section
44
+ with gr.Row():
45
+ with gr.Column():
46
+ # Paper Details View
47
+ paper_view = gr.Markdown(label="Paper Details")
48
+ arxiv_keywords = gr.Textbox(label="Additional Keywords (comma-separated)",
49
+ placeholder="e.g., AI, Deep Learning")
50
+ ingest_button = gr.Button("Ingest Selected Paper")
51
+ ingest_result = gr.Textbox(label="Ingestion Result", interactive=False)
52
+
53
+ # Define States for Pagination and Selection
54
+ state = gr.State(value={"start": 0, "current_page": 1, "last_query": None, "entries": []})
55
+ selected_paper_id = gr.State(value=None)
56
+
57
+ def search_arxiv(query, author, year):
58
+ start = 0
59
+ url = build_query_url(query, author, year, start)
60
+ try:
61
+ response = requests.get(url)
62
+ response.raise_for_status()
63
+ except requests.exceptions.RequestException as e:
64
+ return gr.update(value=[]), gr.update(value=f"**Error:** {str(e)}"), state.value
65
+
66
+ entries = parse_arxiv_feed(response.text)
67
+ state.value = {"start": start, "current_page": 1, "last_query": (query, author, year), "entries": entries}
68
+ if not entries:
69
+ return gr.update(value=[]), "No results found.", state.value
70
+
71
+ # Update the dropdown with paper titles for selection
72
+ titles = [entry['title'] for entry in entries]
73
+ return gr.update(choices=titles), "1", state.value
74
+
75
+ # Dead code? FIXME
76
+ def handle_pagination(direction):
77
+ current_state = state.value
78
+ query, author, year = current_state["last_query"]
79
+ new_page = current_state["current_page"] + direction
80
+ if new_page < 1:
81
+ new_page = 1
82
+ start = (new_page - 1) * ARXIV_PAGE_SIZE
83
+ url = build_query_url(query, author, year, start)
84
+ try:
85
+ response = requests.get(url)
86
+ response.raise_for_status()
87
+ except requests.exceptions.RequestException as e:
88
+ return gr.update(), gr.update()
89
+
90
+ entries = parse_arxiv_feed(response.text)
91
+ if entries:
92
+ current_state["start"] = start
93
+ current_state["current_page"] = new_page
94
+ current_state["entries"] = entries
95
+ state.value = current_state
96
+
97
+ # Update the dropdown with paper titles for the new page
98
+ titles = [entry['title'] for entry in entries]
99
+ return gr.update(choices=titles), str(new_page)
100
+ else:
101
+ # If no entries, do not change the page
102
+ return gr.update(), gr.update()
103
+
104
+ def load_selected_paper(selected_title):
105
+ if not selected_title:
106
+ return "Please select a paper to view."
107
+
108
+ # Find the selected paper from state
109
+ for entry in state.value["entries"]:
110
+ if entry['title'] == selected_title:
111
+ paper_id = entry['id']
112
+ break
113
+ else:
114
+ return "Paper not found."
115
+
116
+ try:
117
+ # Fetch the PDF URL and download the full-text
118
+ pdf_url = fetch_arxiv_pdf_url(paper_id)
119
+ response = requests.get(pdf_url)
120
+ response.raise_for_status()
121
+
122
+ # Save the PDF temporarily
123
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
124
+ temp_pdf.write(response.content)
125
+ temp_pdf_path = temp_pdf.name
126
+
127
+ # Convert PDF to markdown using your PDF ingestion function
128
+ full_text_markdown = extract_text_and_format_from_pdf(temp_pdf_path)
129
+
130
+ selected_paper_id.value = paper_id
131
+ return full_text_markdown
132
+ except Exception as e:
133
+ return f"Error loading full paper: {str(e)}"
134
+
135
+ def process_and_ingest_arxiv_paper(paper_id, additional_keywords):
136
+ try:
137
+ if not paper_id:
138
+ return "Please select a paper to ingest."
139
+
140
+ # Fetch the PDF URL
141
+ pdf_url = fetch_arxiv_pdf_url(paper_id)
142
+
143
+ # Download the PDF
144
+ response = requests.get(pdf_url)
145
+ response.raise_for_status()
146
+
147
+ # Save the PDF temporarily
148
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
149
+ temp_pdf.write(response.content)
150
+ temp_pdf_path = temp_pdf.name
151
+
152
+ # Convert PDF to markdown using your PDF ingestion function
153
+ markdown_text = extract_text_and_format_from_pdf(temp_pdf_path)
154
+
155
+ # Fetch metadata from arXiv to get title, authors, and categories
156
+ xml_content = fetch_arxiv_xml(paper_id)
157
+ _, title, authors, categories = convert_xml_to_markdown(xml_content)
158
+
159
+ # Prepare the arXiv paper URL for access/download
160
+ paper_url = f"https://arxiv.org/abs/{paper_id}"
161
+
162
+ # Prepare the keywords for ingestion
163
+ keywords = f"arxiv,{','.join(categories)}"
164
+ if additional_keywords:
165
+ keywords += f",{additional_keywords}"
166
+
167
+ # Ingest full paper markdown content
168
+ add_media_with_keywords(
169
+ url=paper_url,
170
+ title=title,
171
+ media_type='document',
172
+ content=markdown_text, # Full paper content in markdown
173
+ keywords=keywords,
174
+ prompt='No prompt for arXiv papers',
175
+ summary='Full arXiv paper ingested from PDF',
176
+ transcription_model='None',
177
+ author=', '.join(authors),
178
+ ingestion_date=datetime.now().strftime('%Y-%m-%d')
179
+ )
180
+
181
+ # Return success message with paper title and authors
182
+ return f"arXiv paper '{title}' by {', '.join(authors)} ingested successfully."
183
+ except Exception as e:
184
+ # Return error message if anything goes wrong
185
+ return f"Error processing arXiv paper: {str(e)}"
186
+
187
+ # Event Handlers
188
+ # Connect Search Button
189
+ search_button.click(
190
+ fn=search_arxiv,
191
+ inputs=[search_query, author_filter, year_filter],
192
+ outputs=[paper_selector, page_info, state],
193
+ queue=True
194
+ )
195
+
196
+ # Connect Next Button
197
+ next_button.click(
198
+ fn=lambda: handle_pagination(1),
199
+ inputs=None,
200
+ outputs=[paper_selector, page_info],
201
+ queue=True
202
+ )
203
+
204
+ # Connect Previous Button
205
+ prev_button.click(
206
+ fn=lambda: handle_pagination(-1),
207
+ inputs=None,
208
+ outputs=[paper_selector, page_info],
209
+ queue=True
210
+ )
211
+
212
+ # When the user selects a paper in the Dropdown
213
+ paper_selector.change(
214
+ fn=load_selected_paper,
215
+ inputs=paper_selector,
216
+ outputs=paper_view,
217
+ queue=True
218
+ )
219
+
220
+ # Connect Ingest Button
221
+ ingest_button.click(
222
+ fn=process_and_ingest_arxiv_paper,
223
+ inputs=[selected_paper_id, arxiv_keywords],
224
+ outputs=ingest_result,
225
+ queue=True
226
+ )
227
+
228
+ #
229
+ # End of File
230
+ #####################################################################################################
App_Function_Libraries/Gradio_UI/Audio_ingestion_tab.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Audio_ingestion_tab.py
2
+ # Description: Gradio UI for ingesting audio files into the database
3
+ #
4
+ # Imports
5
+ #
6
+ # External Imports
7
+ import gradio as gr
8
+ #
9
+ # Local Imports
10
+ from App_Function_Libraries.Audio.Audio_Files import process_audio_files
11
+ from App_Function_Libraries.DB.DB_Manager import load_preset_prompts
12
+ from App_Function_Libraries.Gradio_UI.Chat_ui import update_user_prompt
13
+ from App_Function_Libraries.Gradio_UI.Gradio_Shared import whisper_models
14
+ from App_Function_Libraries.Utils.Utils import cleanup_temp_files
15
+ # Import metrics logging
16
+ from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
17
+ from App_Function_Libraries.Metrics.logger_config import logger
18
+ #
19
+ #######################################################################################################################
20
+ # Functions:
21
+
22
+ def create_audio_processing_tab():
23
+ with gr.TabItem("Audio File Transcription + Summarization", visible=True):
24
+ gr.Markdown("# Transcribe & Summarize Audio Files from URLs or Local Files!")
25
+ with gr.Row():
26
+ with gr.Column():
27
+ audio_url_input = gr.Textbox(label="Audio File URL(s)", placeholder="Enter the URL(s) of the audio file(s), one per line")
28
+ audio_file_input = gr.File(label="Upload Audio File", file_types=["audio/*"])
29
+ custom_title_input = gr.Textbox(label="Custom Title/Name", placeholder="Enter a custom title or name for the audio file")
30
+ use_cookies_input = gr.Checkbox(label="Use cookies for authenticated download", value=False)
31
+ cookies_input = gr.Textbox(
32
+ label="Audio Download Cookies",
33
+ placeholder="Paste your cookies here (JSON format)",
34
+ lines=3,
35
+ visible=False
36
+ )
37
+
38
+ use_cookies_input.change(
39
+ fn=lambda x: gr.update(visible=x),
40
+ inputs=[use_cookies_input],
41
+ outputs=[cookies_input]
42
+ )
43
+
44
+ diarize_input = gr.Checkbox(label="Enable Speaker Diarization", value=False)
45
+ whisper_model_input = gr.Dropdown(choices=whisper_models, value="medium", label="Whisper Model")
46
+ keep_timestamps_input = gr.Checkbox(label="Keep Timestamps", value=True)
47
+
48
+ with gr.Row():
49
+ custom_prompt_checkbox = gr.Checkbox(label="Use a Custom Prompt",
50
+ value=False,
51
+ visible=True)
52
+ preset_prompt_checkbox = gr.Checkbox(label="Use a pre-set Prompt",
53
+ value=False,
54
+ visible=True)
55
+ with gr.Row():
56
+ preset_prompt = gr.Dropdown(label="Select Preset Prompt",
57
+ choices=load_preset_prompts(),
58
+ visible=False)
59
+ with gr.Row():
60
+ custom_prompt_input = gr.Textbox(label="Custom Prompt",
61
+ placeholder="Enter custom prompt here",
62
+ lines=3,
63
+ visible=False)
64
+ with gr.Row():
65
+ system_prompt_input = gr.Textbox(label="System Prompt",
66
+ value="""<s>You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST]
67
+ **Bulleted Note Creation Guidelines**
68
+
69
+ **Headings**:
70
+ - Based on referenced topics, not categories like quotes or terms
71
+ - Surrounded by **bold** formatting
72
+ - Not listed as bullet points
73
+ - No space between headings and list items underneath
74
+
75
+ **Emphasis**:
76
+ - **Important terms** set in bold font
77
+ - **Text ending in a colon**: also bolded
78
+
79
+ **Review**:
80
+ - Ensure adherence to specified format
81
+ - Do not reference these instructions in your response.</s>[INST] {{ .Prompt }} [/INST]
82
+ """,
83
+ lines=3,
84
+ visible=False)
85
+
86
+ custom_prompt_checkbox.change(
87
+ fn=lambda x: (gr.update(visible=x), gr.update(visible=x)),
88
+ inputs=[custom_prompt_checkbox],
89
+ outputs=[custom_prompt_input, system_prompt_input]
90
+ )
91
+ preset_prompt_checkbox.change(
92
+ fn=lambda x: gr.update(visible=x),
93
+ inputs=[preset_prompt_checkbox],
94
+ outputs=[preset_prompt]
95
+ )
96
+
97
+ def update_prompts(preset_name):
98
+ prompts = update_user_prompt(preset_name)
99
+ return (
100
+ gr.update(value=prompts["user_prompt"], visible=True),
101
+ gr.update(value=prompts["system_prompt"], visible=True)
102
+ )
103
+
104
+ preset_prompt.change(
105
+ update_prompts,
106
+ inputs=preset_prompt,
107
+ outputs=[custom_prompt_input, system_prompt_input]
108
+ )
109
+
110
+ api_name_input = gr.Dropdown(
111
+ choices=[None, "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", "OpenRouter",
112
+ "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM","ollama", "HuggingFace", "Custom-OpenAI-API"],
113
+ value=None,
114
+ label="API for Summarization (Optional)"
115
+ )
116
+ api_key_input = gr.Textbox(label="API Key (if required)", placeholder="Enter your API key here", type="password")
117
+ custom_keywords_input = gr.Textbox(label="Custom Keywords", placeholder="Enter custom keywords, comma-separated")
118
+ keep_original_input = gr.Checkbox(label="Keep original audio file", value=False)
119
+
120
+ chunking_options_checkbox = gr.Checkbox(label="Show Chunking Options", value=False)
121
+ with gr.Row(visible=False) as chunking_options_box:
122
+ gr.Markdown("### Chunking Options")
123
+ with gr.Column():
124
+ chunk_method = gr.Dropdown(choices=['words', 'sentences', 'paragraphs', 'tokens'], label="Chunking Method")
125
+ max_chunk_size = gr.Slider(minimum=100, maximum=1000, value=300, step=50, label="Max Chunk Size")
126
+ chunk_overlap = gr.Slider(minimum=0, maximum=100, value=0, step=10, label="Chunk Overlap")
127
+ use_adaptive_chunking = gr.Checkbox(label="Use Adaptive Chunking")
128
+ use_multi_level_chunking = gr.Checkbox(label="Use Multi-level Chunking")
129
+ chunk_language = gr.Dropdown(choices=['english', 'french', 'german', 'spanish'], label="Chunking Language")
130
+
131
+ chunking_options_checkbox.change(
132
+ fn=lambda x: gr.update(visible=x),
133
+ inputs=[chunking_options_checkbox],
134
+ outputs=[chunking_options_box]
135
+ )
136
+
137
+ process_audio_button = gr.Button("Process Audio File(s)")
138
+
139
+ with gr.Column():
140
+ audio_progress_output = gr.Textbox(label="Progress")
141
+ audio_transcription_output = gr.Textbox(label="Transcription")
142
+ audio_summary_output = gr.Textbox(label="Summary")
143
+ download_transcription = gr.File(label="Download All Transcriptions as JSON")
144
+ download_summary = gr.File(label="Download All Summaries as Text")
145
+
146
+ process_audio_button.click(
147
+ fn=process_audio_files,
148
+ inputs=[audio_url_input, audio_file_input, whisper_model_input, api_name_input, api_key_input,
149
+ use_cookies_input, cookies_input, keep_original_input, custom_keywords_input, custom_prompt_input,
150
+ chunk_method, max_chunk_size, chunk_overlap, use_adaptive_chunking, use_multi_level_chunking,
151
+ chunk_language, diarize_input, keep_timestamps_input, custom_title_input],
152
+ outputs=[audio_progress_output, audio_transcription_output, audio_summary_output]
153
+ )
154
+
155
+ def on_file_clear(file):
156
+ if file is None:
157
+ cleanup_temp_files()
158
+
159
+ audio_file_input.clear(
160
+ fn=on_file_clear,
161
+ inputs=[audio_file_input],
162
+ outputs=[]
163
+ )
164
+
165
+ #
166
+ # End of Audio_ingestion_tab.py
167
+ #######################################################################################################################