edsl 0.1.39.dev3__py3-none-any.whl → 0.1.39.dev5__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 -332
  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 -867
  7. edsl/agents/AgentList.py +551 -413
  8. edsl/agents/Invigilator.py +284 -233
  9. edsl/agents/InvigilatorBase.py +257 -270
  10. edsl/agents/PromptConstructor.py +272 -354
  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 -157
  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 -1028
  40. edsl/coop/utils.py +131 -131
  41. edsl/data/Cache.py +573 -555
  42. edsl/data/CacheEntry.py +230 -233
  43. edsl/data/CacheHandler.py +168 -149
  44. edsl/data/RemoteCacheSync.py +186 -78
  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 -175
  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 -148
  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 -163
  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 -41
  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 -898
  94. edsl/jobs/JobsChecks.py +172 -147
  95. edsl/jobs/JobsComponentConstructor.py +189 -0
  96. edsl/jobs/JobsPrompts.py +270 -268
  97. edsl/jobs/JobsRemoteInferenceHandler.py +311 -239
  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 -466
  123. edsl/jobs/runners/JobsRunnerStatus.py +297 -330
  124. edsl/jobs/tasks/QuestionTaskCreator.py +244 -242
  125. edsl/jobs/tasks/TaskCreators.py +64 -64
  126. edsl/jobs/tasks/TaskHistory.py +470 -450
  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 -668
  133. edsl/language_models/ModelList.py +164 -155
  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 -362
  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 -664
  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 -182
  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 -177
  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 -108
  220. edsl/results/Dataset.py +587 -424
  221. edsl/results/DatasetExportMixin.py +594 -731
  222. edsl/results/DatasetTree.py +295 -275
  223. edsl/results/MarkdownToDocx.py +122 -0
  224. edsl/results/MarkdownToPDF.py +111 -0
  225. edsl/results/Result.py +557 -465
  226. edsl/results/Results.py +1183 -1165
  227. edsl/results/ResultsExportMixin.py +45 -43
  228. edsl/results/ResultsGGMixin.py +121 -121
  229. edsl/results/TableDisplay.py +125 -198
  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 +77 -77
  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 +511 -632
  245. edsl/scenarios/PdfExtractor.py +40 -0
  246. edsl/scenarios/Scenario.py +498 -601
  247. edsl/scenarios/ScenarioHtmlMixin.py +65 -64
  248. edsl/scenarios/ScenarioList.py +1458 -1287
  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 +38 -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/{ScenarioJoin.py → scenario_join.py} +131 -127
  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 -1801
  288. edsl/surveys/SurveyCSS.py +273 -261
  289. edsl/surveys/SurveyExportMixin.py +259 -259
  290. edsl/surveys/{SurveyFlowVisualizationMixin.py → SurveyFlowVisualization.py} +181 -179
  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 -65
  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 -19
  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 -424
  333. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/LICENSE +21 -21
  334. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/METADATA +13 -11
  335. edsl-0.1.39.dev5.dist-info/RECORD +358 -0
  336. {edsl-0.1.39.dev3.dist-info → edsl-0.1.39.dev5.dist-info}/WHEEL +1 -1
  337. edsl/language_models/KeyLookup.py +0 -30
  338. edsl/language_models/registry.py +0 -190
  339. edsl/language_models/unused/ReplicateBase.py +0 -83
  340. edsl/results/ResultsDBMixin.py +0 -238
  341. edsl-0.1.39.dev3.dist-info/RECORD +0 -277
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