edsl 0.1.38.dev3__py3-none-any.whl → 0.1.39__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (341) hide show
  1. edsl/Base.py +413 -303
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +57 -49
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +1071 -858
  7. edsl/agents/AgentList.py +551 -362
  8. edsl/agents/Invigilator.py +284 -222
  9. edsl/agents/InvigilatorBase.py +257 -284
  10. edsl/agents/PromptConstructor.py +272 -353
  11. edsl/agents/QuestionInstructionPromptBuilder.py +128 -0
  12. edsl/agents/QuestionTemplateReplacementsBuilder.py +137 -0
  13. edsl/agents/__init__.py +2 -3
  14. edsl/agents/descriptors.py +99 -99
  15. edsl/agents/prompt_helpers.py +129 -129
  16. edsl/agents/question_option_processor.py +172 -0
  17. edsl/auto/AutoStudy.py +130 -117
  18. edsl/auto/StageBase.py +243 -230
  19. edsl/auto/StageGenerateSurvey.py +178 -178
  20. edsl/auto/StageLabelQuestions.py +125 -125
  21. edsl/auto/StagePersona.py +61 -61
  22. edsl/auto/StagePersonaDimensionValueRanges.py +88 -88
  23. edsl/auto/StagePersonaDimensionValues.py +74 -74
  24. edsl/auto/StagePersonaDimensions.py +69 -69
  25. edsl/auto/StageQuestions.py +74 -73
  26. edsl/auto/SurveyCreatorPipeline.py +21 -21
  27. edsl/auto/utilities.py +218 -224
  28. edsl/base/Base.py +279 -279
  29. edsl/config.py +177 -149
  30. edsl/conversation/Conversation.py +290 -290
  31. edsl/conversation/car_buying.py +59 -58
  32. edsl/conversation/chips.py +95 -95
  33. edsl/conversation/mug_negotiation.py +81 -81
  34. edsl/conversation/next_speaker_utilities.py +93 -93
  35. edsl/coop/CoopFunctionsMixin.py +15 -0
  36. edsl/coop/ExpectedParrotKeyHandler.py +125 -0
  37. edsl/coop/PriceFetcher.py +54 -54
  38. edsl/coop/__init__.py +2 -2
  39. edsl/coop/coop.py +1106 -961
  40. edsl/coop/utils.py +131 -131
  41. edsl/data/Cache.py +573 -530
  42. edsl/data/CacheEntry.py +230 -228
  43. edsl/data/CacheHandler.py +168 -149
  44. edsl/data/RemoteCacheSync.py +186 -97
  45. edsl/data/SQLiteDict.py +292 -292
  46. edsl/data/__init__.py +5 -4
  47. edsl/data/orm.py +10 -10
  48. edsl/data_transfer_models.py +74 -73
  49. edsl/enums.py +202 -173
  50. edsl/exceptions/BaseException.py +21 -21
  51. edsl/exceptions/__init__.py +54 -54
  52. edsl/exceptions/agents.py +54 -42
  53. edsl/exceptions/cache.py +5 -5
  54. edsl/exceptions/configuration.py +16 -16
  55. edsl/exceptions/coop.py +10 -10
  56. edsl/exceptions/data.py +14 -14
  57. edsl/exceptions/general.py +34 -34
  58. edsl/exceptions/inference_services.py +5 -0
  59. edsl/exceptions/jobs.py +33 -33
  60. edsl/exceptions/language_models.py +63 -63
  61. edsl/exceptions/prompts.py +15 -15
  62. edsl/exceptions/questions.py +109 -91
  63. edsl/exceptions/results.py +29 -29
  64. edsl/exceptions/scenarios.py +29 -22
  65. edsl/exceptions/surveys.py +37 -37
  66. edsl/inference_services/AnthropicService.py +106 -87
  67. edsl/inference_services/AvailableModelCacheHandler.py +184 -0
  68. edsl/inference_services/AvailableModelFetcher.py +215 -0
  69. edsl/inference_services/AwsBedrock.py +118 -120
  70. edsl/inference_services/AzureAI.py +215 -217
  71. edsl/inference_services/DeepInfraService.py +18 -18
  72. edsl/inference_services/GoogleService.py +143 -156
  73. edsl/inference_services/GroqService.py +20 -20
  74. edsl/inference_services/InferenceServiceABC.py +80 -147
  75. edsl/inference_services/InferenceServicesCollection.py +138 -97
  76. edsl/inference_services/MistralAIService.py +120 -123
  77. edsl/inference_services/OllamaService.py +18 -18
  78. edsl/inference_services/OpenAIService.py +236 -224
  79. edsl/inference_services/PerplexityService.py +160 -0
  80. edsl/inference_services/ServiceAvailability.py +135 -0
  81. edsl/inference_services/TestService.py +90 -89
  82. edsl/inference_services/TogetherAIService.py +172 -170
  83. edsl/inference_services/data_structures.py +134 -0
  84. edsl/inference_services/models_available_cache.py +118 -118
  85. edsl/inference_services/rate_limits_cache.py +25 -25
  86. edsl/inference_services/registry.py +41 -39
  87. edsl/inference_services/write_available.py +10 -10
  88. edsl/jobs/AnswerQuestionFunctionConstructor.py +223 -0
  89. edsl/jobs/Answers.py +43 -56
  90. edsl/jobs/FetchInvigilator.py +47 -0
  91. edsl/jobs/InterviewTaskManager.py +98 -0
  92. edsl/jobs/InterviewsConstructor.py +50 -0
  93. edsl/jobs/Jobs.py +823 -1358
  94. edsl/jobs/JobsChecks.py +172 -0
  95. edsl/jobs/JobsComponentConstructor.py +189 -0
  96. edsl/jobs/JobsPrompts.py +270 -0
  97. edsl/jobs/JobsRemoteInferenceHandler.py +311 -0
  98. edsl/jobs/JobsRemoteInferenceLogger.py +239 -0
  99. edsl/jobs/RequestTokenEstimator.py +30 -0
  100. edsl/jobs/__init__.py +1 -1
  101. edsl/jobs/async_interview_runner.py +138 -0
  102. edsl/jobs/buckets/BucketCollection.py +104 -63
  103. edsl/jobs/buckets/ModelBuckets.py +65 -65
  104. edsl/jobs/buckets/TokenBucket.py +283 -251
  105. edsl/jobs/buckets/TokenBucketAPI.py +211 -0
  106. edsl/jobs/buckets/TokenBucketClient.py +191 -0
  107. edsl/jobs/check_survey_scenario_compatibility.py +85 -0
  108. edsl/jobs/data_structures.py +120 -0
  109. edsl/jobs/decorators.py +35 -0
  110. edsl/jobs/interviews/Interview.py +396 -661
  111. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  112. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  113. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  114. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  115. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  116. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  117. edsl/jobs/interviews/ReportErrors.py +66 -66
  118. edsl/jobs/interviews/interview_status_enum.py +9 -9
  119. edsl/jobs/jobs_status_enums.py +9 -0
  120. edsl/jobs/loggers/HTMLTableJobLogger.py +304 -0
  121. edsl/jobs/results_exceptions_handler.py +98 -0
  122. edsl/jobs/runners/JobsRunnerAsyncio.py +151 -361
  123. edsl/jobs/runners/JobsRunnerStatus.py +298 -332
  124. edsl/jobs/tasks/QuestionTaskCreator.py +244 -242
  125. edsl/jobs/tasks/TaskCreators.py +64 -64
  126. edsl/jobs/tasks/TaskHistory.py +470 -451
  127. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  128. edsl/jobs/tasks/task_status_enum.py +161 -163
  129. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  130. edsl/jobs/tokens/TokenUsage.py +34 -34
  131. edsl/language_models/ComputeCost.py +63 -0
  132. edsl/language_models/LanguageModel.py +626 -708
  133. edsl/language_models/ModelList.py +164 -109
  134. edsl/language_models/PriceManager.py +127 -0
  135. edsl/language_models/RawResponseHandler.py +106 -0
  136. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  137. edsl/language_models/ServiceDataSources.py +0 -0
  138. edsl/language_models/__init__.py +2 -3
  139. edsl/language_models/fake_openai_call.py +15 -15
  140. edsl/language_models/fake_openai_service.py +61 -61
  141. edsl/language_models/key_management/KeyLookup.py +63 -0
  142. edsl/language_models/key_management/KeyLookupBuilder.py +273 -0
  143. edsl/language_models/key_management/KeyLookupCollection.py +38 -0
  144. edsl/language_models/key_management/__init__.py +0 -0
  145. edsl/language_models/key_management/models.py +131 -0
  146. edsl/language_models/model.py +256 -0
  147. edsl/language_models/repair.py +156 -156
  148. edsl/language_models/utilities.py +65 -64
  149. edsl/notebooks/Notebook.py +263 -258
  150. edsl/notebooks/NotebookToLaTeX.py +142 -0
  151. edsl/notebooks/__init__.py +1 -1
  152. edsl/prompts/Prompt.py +352 -357
  153. edsl/prompts/__init__.py +2 -2
  154. edsl/questions/ExceptionExplainer.py +77 -0
  155. edsl/questions/HTMLQuestion.py +103 -0
  156. edsl/questions/QuestionBase.py +518 -660
  157. edsl/questions/QuestionBasePromptsMixin.py +221 -217
  158. edsl/questions/QuestionBudget.py +227 -227
  159. edsl/questions/QuestionCheckBox.py +359 -359
  160. edsl/questions/QuestionExtract.py +180 -183
  161. edsl/questions/QuestionFreeText.py +113 -114
  162. edsl/questions/QuestionFunctional.py +166 -166
  163. edsl/questions/QuestionList.py +223 -231
  164. edsl/questions/QuestionMatrix.py +265 -0
  165. edsl/questions/QuestionMultipleChoice.py +330 -286
  166. edsl/questions/QuestionNumerical.py +151 -153
  167. edsl/questions/QuestionRank.py +314 -324
  168. edsl/questions/Quick.py +41 -41
  169. edsl/questions/SimpleAskMixin.py +74 -73
  170. edsl/questions/__init__.py +27 -26
  171. edsl/questions/{AnswerValidatorMixin.py → answer_validator_mixin.py} +334 -289
  172. edsl/questions/compose_questions.py +98 -98
  173. edsl/questions/data_structures.py +20 -0
  174. edsl/questions/decorators.py +21 -21
  175. edsl/questions/derived/QuestionLikertFive.py +76 -76
  176. edsl/questions/derived/QuestionLinearScale.py +90 -87
  177. edsl/questions/derived/QuestionTopK.py +93 -93
  178. edsl/questions/derived/QuestionYesNo.py +82 -82
  179. edsl/questions/descriptors.py +427 -413
  180. edsl/questions/loop_processor.py +149 -0
  181. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  182. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  183. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  184. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  185. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  186. edsl/questions/prompt_templates/question_list.jinja +17 -17
  187. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  188. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  189. edsl/questions/{QuestionBaseGenMixin.py → question_base_gen_mixin.py} +168 -161
  190. edsl/questions/question_registry.py +177 -147
  191. edsl/questions/{RegisterQuestionsMeta.py → register_questions_meta.py} +71 -71
  192. edsl/questions/{ResponseValidatorABC.py → response_validator_abc.py} +188 -174
  193. edsl/questions/response_validator_factory.py +34 -0
  194. edsl/questions/settings.py +12 -12
  195. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  196. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  197. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  198. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  199. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  200. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  201. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  202. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  203. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  204. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  205. edsl/questions/templates/list/question_presentation.jinja +5 -5
  206. edsl/questions/templates/matrix/__init__.py +1 -0
  207. edsl/questions/templates/matrix/answering_instructions.jinja +5 -0
  208. edsl/questions/templates/matrix/question_presentation.jinja +20 -0
  209. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  210. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  211. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  212. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  213. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  214. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  215. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  216. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  217. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  218. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  219. edsl/results/CSSParameterizer.py +108 -0
  220. edsl/results/Dataset.py +587 -293
  221. edsl/results/DatasetExportMixin.py +594 -717
  222. edsl/results/DatasetTree.py +295 -145
  223. edsl/results/MarkdownToDocx.py +122 -0
  224. edsl/results/MarkdownToPDF.py +111 -0
  225. edsl/results/Result.py +557 -456
  226. edsl/results/Results.py +1183 -1071
  227. edsl/results/ResultsExportMixin.py +45 -43
  228. edsl/results/ResultsGGMixin.py +121 -121
  229. edsl/results/TableDisplay.py +125 -0
  230. edsl/results/TextEditor.py +50 -0
  231. edsl/results/__init__.py +2 -2
  232. edsl/results/file_exports.py +252 -0
  233. edsl/results/{ResultsFetchMixin.py → results_fetch_mixin.py} +33 -33
  234. edsl/results/{Selector.py → results_selector.py} +145 -135
  235. edsl/results/{ResultsToolsMixin.py → results_tools_mixin.py} +98 -98
  236. edsl/results/smart_objects.py +96 -0
  237. edsl/results/table_data_class.py +12 -0
  238. edsl/results/table_display.css +78 -0
  239. edsl/results/table_renderers.py +118 -0
  240. edsl/results/tree_explore.py +115 -115
  241. edsl/scenarios/ConstructDownloadLink.py +109 -0
  242. edsl/scenarios/DocumentChunker.py +102 -0
  243. edsl/scenarios/DocxScenario.py +16 -0
  244. edsl/scenarios/FileStore.py +543 -458
  245. edsl/scenarios/PdfExtractor.py +40 -0
  246. edsl/scenarios/Scenario.py +498 -544
  247. edsl/scenarios/ScenarioHtmlMixin.py +65 -64
  248. edsl/scenarios/ScenarioList.py +1458 -1112
  249. edsl/scenarios/ScenarioListExportMixin.py +45 -52
  250. edsl/scenarios/ScenarioListPdfMixin.py +239 -261
  251. edsl/scenarios/__init__.py +3 -4
  252. edsl/scenarios/directory_scanner.py +96 -0
  253. edsl/scenarios/file_methods.py +85 -0
  254. edsl/scenarios/handlers/__init__.py +13 -0
  255. edsl/scenarios/handlers/csv.py +49 -0
  256. edsl/scenarios/handlers/docx.py +76 -0
  257. edsl/scenarios/handlers/html.py +37 -0
  258. edsl/scenarios/handlers/json.py +111 -0
  259. edsl/scenarios/handlers/latex.py +5 -0
  260. edsl/scenarios/handlers/md.py +51 -0
  261. edsl/scenarios/handlers/pdf.py +68 -0
  262. edsl/scenarios/handlers/png.py +39 -0
  263. edsl/scenarios/handlers/pptx.py +105 -0
  264. edsl/scenarios/handlers/py.py +294 -0
  265. edsl/scenarios/handlers/sql.py +313 -0
  266. edsl/scenarios/handlers/sqlite.py +149 -0
  267. edsl/scenarios/handlers/txt.py +33 -0
  268. edsl/scenarios/scenario_join.py +131 -0
  269. edsl/scenarios/scenario_selector.py +156 -0
  270. edsl/shared.py +1 -1
  271. edsl/study/ObjectEntry.py +173 -173
  272. edsl/study/ProofOfWork.py +113 -113
  273. edsl/study/SnapShot.py +80 -80
  274. edsl/study/Study.py +521 -528
  275. edsl/study/__init__.py +4 -4
  276. edsl/surveys/ConstructDAG.py +92 -0
  277. edsl/surveys/DAG.py +148 -148
  278. edsl/surveys/EditSurvey.py +221 -0
  279. edsl/surveys/InstructionHandler.py +100 -0
  280. edsl/surveys/Memory.py +31 -31
  281. edsl/surveys/MemoryManagement.py +72 -0
  282. edsl/surveys/MemoryPlan.py +244 -244
  283. edsl/surveys/Rule.py +327 -326
  284. edsl/surveys/RuleCollection.py +385 -387
  285. edsl/surveys/RuleManager.py +172 -0
  286. edsl/surveys/Simulator.py +75 -0
  287. edsl/surveys/Survey.py +1280 -1787
  288. edsl/surveys/SurveyCSS.py +273 -261
  289. edsl/surveys/SurveyExportMixin.py +259 -259
  290. edsl/surveys/{SurveyFlowVisualizationMixin.py → SurveyFlowVisualization.py} +181 -121
  291. edsl/surveys/SurveyQualtricsImport.py +284 -284
  292. edsl/surveys/SurveyToApp.py +141 -0
  293. edsl/surveys/__init__.py +5 -3
  294. edsl/surveys/base.py +53 -53
  295. edsl/surveys/descriptors.py +60 -56
  296. edsl/surveys/instructions/ChangeInstruction.py +48 -49
  297. edsl/surveys/instructions/Instruction.py +56 -53
  298. edsl/surveys/instructions/InstructionCollection.py +82 -77
  299. edsl/templates/error_reporting/base.html +23 -23
  300. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  301. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  302. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  303. edsl/templates/error_reporting/interview_details.html +115 -115
  304. edsl/templates/error_reporting/interviews.html +19 -10
  305. edsl/templates/error_reporting/overview.html +4 -4
  306. edsl/templates/error_reporting/performance_plot.html +1 -1
  307. edsl/templates/error_reporting/report.css +73 -73
  308. edsl/templates/error_reporting/report.html +117 -117
  309. edsl/templates/error_reporting/report.js +25 -25
  310. edsl/tools/__init__.py +1 -1
  311. edsl/tools/clusters.py +192 -192
  312. edsl/tools/embeddings.py +27 -27
  313. edsl/tools/embeddings_plotting.py +118 -118
  314. edsl/tools/plotting.py +112 -112
  315. edsl/tools/summarize.py +18 -18
  316. edsl/utilities/PrettyList.py +56 -0
  317. edsl/utilities/SystemInfo.py +28 -28
  318. edsl/utilities/__init__.py +22 -22
  319. edsl/utilities/ast_utilities.py +25 -25
  320. edsl/utilities/data/Registry.py +6 -6
  321. edsl/utilities/data/__init__.py +1 -1
  322. edsl/utilities/data/scooter_results.json +1 -1
  323. edsl/utilities/decorators.py +77 -77
  324. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  325. edsl/utilities/interface.py +627 -627
  326. edsl/utilities/is_notebook.py +18 -0
  327. edsl/utilities/is_valid_variable_name.py +11 -0
  328. edsl/utilities/naming_utilities.py +263 -263
  329. edsl/utilities/remove_edsl_version.py +24 -0
  330. edsl/utilities/repair_functions.py +28 -28
  331. edsl/utilities/restricted_python.py +70 -70
  332. edsl/utilities/utilities.py +436 -409
  333. {edsl-0.1.38.dev3.dist-info → edsl-0.1.39.dist-info}/LICENSE +21 -21
  334. {edsl-0.1.38.dev3.dist-info → edsl-0.1.39.dist-info}/METADATA +13 -10
  335. edsl-0.1.39.dist-info/RECORD +358 -0
  336. {edsl-0.1.38.dev3.dist-info → edsl-0.1.39.dist-info}/WHEEL +1 -1
  337. edsl/language_models/KeyLookup.py +0 -30
  338. edsl/language_models/registry.py +0 -137
  339. edsl/language_models/unused/ReplicateBase.py +0 -83
  340. edsl/results/ResultsDBMixin.py +0 -238
  341. edsl-0.1.38.dev3.dist-info/RECORD +0 -269
edsl/tools/clusters.py CHANGED
@@ -1,192 +1,192 @@
1
- import json
2
- import numpy as np
3
- from sklearn.cluster import KMeans
4
- from sklearn.manifold import TSNE
5
- from sklearn.decomposition import PCA
6
- from IPython.display import display_html
7
-
8
-
9
- def compute_tsne(embeddings, cluster_labels, text_labels):
10
- """
11
- Compute t-SNE on embedding vectors.
12
-
13
- Parameters:
14
- embeddings (np.ndarray): The embedding vectors.
15
- cluster_labels (np.ndarray): Cluster labels for each embedding.
16
- text_labels (list): Text labels for each embedding.
17
-
18
- Returns:
19
- list: List of dictionaries with x, y coordinates, cluster labels, and text labels.
20
- """
21
- tsne = TSNE(n_components=2, random_state=42)
22
- tsne_results = tsne.fit_transform(embeddings)
23
- data = [
24
- {
25
- "x": float(tsne_results[i, 0]),
26
- "y": float(tsne_results[i, 1]),
27
- "cluster_label": str(cluster_labels[i]),
28
- "text_label": text_labels[i],
29
- }
30
- for i in range(len(cluster_labels))
31
- ]
32
- return data
33
-
34
-
35
- def compute_pca(embeddings, cluster_labels, text_labels):
36
- """
37
- Compute PCA on embedding vectors.
38
-
39
- Parameters:
40
- embeddings (np.ndarray): The embedding vectors.
41
- cluster_labels (np.ndarray): Cluster labels for each embedding.
42
- text_labels (list): Text labels for each embedding.
43
-
44
- Returns:
45
- list: List of dictionaries with x, y coordinates, cluster labels, and text labels.
46
- """
47
- pca = PCA(n_components=2)
48
- pca_results = pca.fit_transform(embeddings)
49
- data = [
50
- {
51
- "x": float(pca_results[i, 0]),
52
- "y": float(pca_results[i, 1]),
53
- "cluster_label": str(cluster_labels[i]),
54
- "text_label": text_labels[i],
55
- }
56
- for i in range(len(cluster_labels))
57
- ]
58
- return data
59
-
60
-
61
- def plot(embeddings, text_labels, n_clusters=5, method="tsne"):
62
- """
63
- Perform k-means clustering and plot results in a Jupyter notebook using D3.js.
64
-
65
- Parameters:
66
- embeddings (np.ndarray): The embedding vectors.
67
- text_labels (list): Text labels for each embedding.
68
- n_clusters (int): The number of clusters to form.
69
- method (str): The dimensionality reduction method to use ('tsne' or 'pca').
70
- """
71
- # Perform k-means clustering
72
- kmeans = KMeans(n_clusters=n_clusters, random_state=42)
73
- cluster_labels = kmeans.fit_predict(embeddings)
74
-
75
- # Compute dimensionality reduction
76
- if method == "tsne":
77
- data = compute_tsne(embeddings, cluster_labels, text_labels)
78
- elif method == "pca":
79
- data = compute_pca(embeddings, cluster_labels, text_labels)
80
- else:
81
- raise ValueError("Invalid method. Choose 'tsne' or 'pca'.")
82
-
83
- # Convert data to JSON
84
- data_json = json.dumps(data)
85
-
86
- # HTML content with embedded data
87
- html_content = f"""
88
- <!DOCTYPE html>
89
- <html lang="en">
90
- <head>
91
- <meta charset="UTF-8">
92
- <title>t-SNE/PCA Plot with D3.js</title>
93
- <script src="https://d3js.org/d3.v6.min.js"></script>
94
- <style>
95
- .tooltip {{
96
- position: absolute;
97
- text-align: center;
98
- width: auto;
99
- height: auto;
100
- padding: 2px;
101
- font: 12px sans-serif;
102
- background: lightsteelblue;
103
- border: 0px;
104
- border-radius: 8px;
105
- pointer-events: none;
106
- }}
107
- .dot {{
108
- stroke: #000;
109
- stroke-width: 0.5;
110
- }}
111
- </style>
112
- </head>
113
- <body>
114
- <svg width="600" height="600"></svg>
115
-
116
- <script>
117
- // Embedded data
118
- const data = {data_json};
119
-
120
- const svg = d3.select("svg"),
121
- width = +svg.attr("width"),
122
- height = +svg.attr("height");
123
-
124
- // Set up scales
125
- const x = d3.scaleLinear()
126
- .domain(d3.extent(data, d => d.x))
127
- .range([0, width]);
128
-
129
- const y = d3.scaleLinear()
130
- .domain(d3.extent(data, d => d.y))
131
- .range([height, 0]);
132
-
133
- // Set up color scale
134
- const color = d3.scaleOrdinal(d3.schemeCategory10);
135
-
136
- // Create tooltip
137
- const tooltip = d3.select("body").append("div")
138
- .attr("class", "tooltip")
139
- .style("opacity", 0);
140
-
141
- // Create circles for each point
142
- svg.selectAll("circle")
143
- .data(data)
144
- .enter().append("circle")
145
- .attr("cx", d => x(d.x))
146
- .attr("cy", d => y(d.y))
147
- .attr("r", 5)
148
- .attr("class", "dot")
149
- .style("fill", d => color(d.cluster_label))
150
- .on("mouseover", function(event, d) {{
151
- tooltip.transition()
152
- .duration(200)
153
- .style("opacity", .9);
154
- tooltip.html(d.text_label)
155
- .style("left", (event.pageX + 5) + "px")
156
- .style("top", (event.pageY - 28) + "px");
157
- }})
158
- .on("mouseout", function(d) {{
159
- tooltip.transition()
160
- .duration(500)
161
- .style("opacity", 0);
162
- }});
163
- </script>
164
- </body>
165
- </html>
166
- """
167
-
168
- # Write HTML content to a temporary file
169
- html_file = "tsne_pca_plot.html"
170
- with open(html_file, "w") as file:
171
- file.write(html_content)
172
-
173
- # Display the HTML content in an iframe within a Jupyter notebook
174
- display_html(
175
- f'<iframe src="{html_file}" width="600" height="600"></iframe>', raw=True
176
- )
177
-
178
-
179
- # Example usage
180
- if __name__ == "__main__":
181
- # Generate some sample data (embedding vectors)
182
- np.random.seed(42)
183
- embedding_vectors = np.random.rand(
184
- 100, 50
185
- ) # 100 samples with 50-dimensional embeddings
186
- text_labels = [f"Text {i}" for i in range(100)] # Sample text labels
187
-
188
- # Plot the clusters using t-SNE
189
- plot(embedding_vectors, text_labels, n_clusters=5, method="tsne")
190
-
191
- # Plot the clusters using PCA
192
- plot(embedding_vectors, text_labels, n_clusters=5, method="pca")
1
+ import json
2
+ import numpy as np
3
+ from sklearn.cluster import KMeans
4
+ from sklearn.manifold import TSNE
5
+ from sklearn.decomposition import PCA
6
+ from IPython.display import display_html
7
+
8
+
9
+ def compute_tsne(embeddings, cluster_labels, text_labels):
10
+ """
11
+ Compute t-SNE on embedding vectors.
12
+
13
+ Parameters:
14
+ embeddings (np.ndarray): The embedding vectors.
15
+ cluster_labels (np.ndarray): Cluster labels for each embedding.
16
+ text_labels (list): Text labels for each embedding.
17
+
18
+ Returns:
19
+ list: List of dictionaries with x, y coordinates, cluster labels, and text labels.
20
+ """
21
+ tsne = TSNE(n_components=2, random_state=42)
22
+ tsne_results = tsne.fit_transform(embeddings)
23
+ data = [
24
+ {
25
+ "x": float(tsne_results[i, 0]),
26
+ "y": float(tsne_results[i, 1]),
27
+ "cluster_label": str(cluster_labels[i]),
28
+ "text_label": text_labels[i],
29
+ }
30
+ for i in range(len(cluster_labels))
31
+ ]
32
+ return data
33
+
34
+
35
+ def compute_pca(embeddings, cluster_labels, text_labels):
36
+ """
37
+ Compute PCA on embedding vectors.
38
+
39
+ Parameters:
40
+ embeddings (np.ndarray): The embedding vectors.
41
+ cluster_labels (np.ndarray): Cluster labels for each embedding.
42
+ text_labels (list): Text labels for each embedding.
43
+
44
+ Returns:
45
+ list: List of dictionaries with x, y coordinates, cluster labels, and text labels.
46
+ """
47
+ pca = PCA(n_components=2)
48
+ pca_results = pca.fit_transform(embeddings)
49
+ data = [
50
+ {
51
+ "x": float(pca_results[i, 0]),
52
+ "y": float(pca_results[i, 1]),
53
+ "cluster_label": str(cluster_labels[i]),
54
+ "text_label": text_labels[i],
55
+ }
56
+ for i in range(len(cluster_labels))
57
+ ]
58
+ return data
59
+
60
+
61
+ def plot(embeddings, text_labels, n_clusters=5, method="tsne"):
62
+ """
63
+ Perform k-means clustering and plot results in a Jupyter notebook using D3.js.
64
+
65
+ Parameters:
66
+ embeddings (np.ndarray): The embedding vectors.
67
+ text_labels (list): Text labels for each embedding.
68
+ n_clusters (int): The number of clusters to form.
69
+ method (str): The dimensionality reduction method to use ('tsne' or 'pca').
70
+ """
71
+ # Perform k-means clustering
72
+ kmeans = KMeans(n_clusters=n_clusters, random_state=42)
73
+ cluster_labels = kmeans.fit_predict(embeddings)
74
+
75
+ # Compute dimensionality reduction
76
+ if method == "tsne":
77
+ data = compute_tsne(embeddings, cluster_labels, text_labels)
78
+ elif method == "pca":
79
+ data = compute_pca(embeddings, cluster_labels, text_labels)
80
+ else:
81
+ raise ValueError("Invalid method. Choose 'tsne' or 'pca'.")
82
+
83
+ # Convert data to JSON
84
+ data_json = json.dumps(data)
85
+
86
+ # HTML content with embedded data
87
+ html_content = f"""
88
+ <!DOCTYPE html>
89
+ <html lang="en">
90
+ <head>
91
+ <meta charset="UTF-8">
92
+ <title>t-SNE/PCA Plot with D3.js</title>
93
+ <script src="https://d3js.org/d3.v6.min.js"></script>
94
+ <style>
95
+ .tooltip {{
96
+ position: absolute;
97
+ text-align: center;
98
+ width: auto;
99
+ height: auto;
100
+ padding: 2px;
101
+ font: 12px sans-serif;
102
+ background: lightsteelblue;
103
+ border: 0px;
104
+ border-radius: 8px;
105
+ pointer-events: none;
106
+ }}
107
+ .dot {{
108
+ stroke: #000;
109
+ stroke-width: 0.5;
110
+ }}
111
+ </style>
112
+ </head>
113
+ <body>
114
+ <svg width="600" height="600"></svg>
115
+
116
+ <script>
117
+ // Embedded data
118
+ const data = {data_json};
119
+
120
+ const svg = d3.select("svg"),
121
+ width = +svg.attr("width"),
122
+ height = +svg.attr("height");
123
+
124
+ // Set up scales
125
+ const x = d3.scaleLinear()
126
+ .domain(d3.extent(data, d => d.x))
127
+ .range([0, width]);
128
+
129
+ const y = d3.scaleLinear()
130
+ .domain(d3.extent(data, d => d.y))
131
+ .range([height, 0]);
132
+
133
+ // Set up color scale
134
+ const color = d3.scaleOrdinal(d3.schemeCategory10);
135
+
136
+ // Create tooltip
137
+ const tooltip = d3.select("body").append("div")
138
+ .attr("class", "tooltip")
139
+ .style("opacity", 0);
140
+
141
+ // Create circles for each point
142
+ svg.selectAll("circle")
143
+ .data(data)
144
+ .enter().append("circle")
145
+ .attr("cx", d => x(d.x))
146
+ .attr("cy", d => y(d.y))
147
+ .attr("r", 5)
148
+ .attr("class", "dot")
149
+ .style("fill", d => color(d.cluster_label))
150
+ .on("mouseover", function(event, d) {{
151
+ tooltip.transition()
152
+ .duration(200)
153
+ .style("opacity", .9);
154
+ tooltip.html(d.text_label)
155
+ .style("left", (event.pageX + 5) + "px")
156
+ .style("top", (event.pageY - 28) + "px");
157
+ }})
158
+ .on("mouseout", function(d) {{
159
+ tooltip.transition()
160
+ .duration(500)
161
+ .style("opacity", 0);
162
+ }});
163
+ </script>
164
+ </body>
165
+ </html>
166
+ """
167
+
168
+ # Write HTML content to a temporary file
169
+ html_file = "tsne_pca_plot.html"
170
+ with open(html_file, "w") as file:
171
+ file.write(html_content)
172
+
173
+ # Display the HTML content in an iframe within a Jupyter notebook
174
+ display_html(
175
+ f'<iframe src="{html_file}" width="600" height="600"></iframe>', raw=True
176
+ )
177
+
178
+
179
+ # Example usage
180
+ if __name__ == "__main__":
181
+ # Generate some sample data (embedding vectors)
182
+ np.random.seed(42)
183
+ embedding_vectors = np.random.rand(
184
+ 100, 50
185
+ ) # 100 samples with 50-dimensional embeddings
186
+ text_labels = [f"Text {i}" for i in range(100)] # Sample text labels
187
+
188
+ # Plot the clusters using t-SNE
189
+ plot(embedding_vectors, text_labels, n_clusters=5, method="tsne")
190
+
191
+ # Plot the clusters using PCA
192
+ plot(embedding_vectors, text_labels, n_clusters=5, method="pca")
edsl/tools/embeddings.py CHANGED
@@ -1,27 +1,27 @@
1
- import openai
2
- from dotenv import load_dotenv
3
- import os
4
-
5
- # Load environment variables from .env file
6
- load_dotenv()
7
-
8
- # Get the OpenAI API key from the environment variable
9
- openai.api_key = os.getenv("OPENAI_API_KEY")
10
-
11
- from openai import Client
12
-
13
-
14
- def get_embeddings(texts):
15
- """
16
- Get embeddings for a list of texts using OpenAI API.
17
-
18
- Args:
19
- texts (list of str): List of strings to get embeddings for.
20
-
21
- Returns:
22
- list of list of float: List of embeddings.
23
- """
24
- client = Client()
25
- response = client.embeddings.create(input=texts, model="text-embedding-ada-002")
26
- embeddings = [item.embedding for item in response.data]
27
- return embeddings
1
+ import openai
2
+ from dotenv import load_dotenv
3
+ import os
4
+
5
+ # Load environment variables from .env file
6
+ load_dotenv()
7
+
8
+ # Get the OpenAI API key from the environment variable
9
+ openai.api_key = os.getenv("OPENAI_API_KEY")
10
+
11
+ from openai import Client
12
+
13
+
14
+ def get_embeddings(texts):
15
+ """
16
+ Get embeddings for a list of texts using OpenAI API.
17
+
18
+ Args:
19
+ texts (list of str): List of strings to get embeddings for.
20
+
21
+ Returns:
22
+ list of list of float: List of embeddings.
23
+ """
24
+ client = Client()
25
+ response = client.embeddings.create(input=texts, model="text-embedding-ada-002")
26
+ embeddings = [item.embedding for item in response.data]
27
+ return embeddings