edsl 0.1.37.dev6__py3-none-any.whl → 0.1.38__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 (271) hide show
  1. edsl/Base.py +332 -303
  2. edsl/BaseDiff.py +260 -260
  3. edsl/TemplateLoader.py +24 -24
  4. edsl/__init__.py +49 -48
  5. edsl/__version__.py +1 -1
  6. edsl/agents/Agent.py +867 -855
  7. edsl/agents/AgentList.py +413 -350
  8. edsl/agents/Invigilator.py +233 -222
  9. edsl/agents/InvigilatorBase.py +265 -284
  10. edsl/agents/PromptConstructor.py +354 -353
  11. edsl/agents/__init__.py +3 -3
  12. edsl/agents/descriptors.py +99 -99
  13. edsl/agents/prompt_helpers.py +129 -129
  14. edsl/auto/AutoStudy.py +117 -117
  15. edsl/auto/StageBase.py +230 -230
  16. edsl/auto/StageGenerateSurvey.py +178 -178
  17. edsl/auto/StageLabelQuestions.py +125 -125
  18. edsl/auto/StagePersona.py +61 -61
  19. edsl/auto/StagePersonaDimensionValueRanges.py +88 -88
  20. edsl/auto/StagePersonaDimensionValues.py +74 -74
  21. edsl/auto/StagePersonaDimensions.py +69 -69
  22. edsl/auto/StageQuestions.py +73 -73
  23. edsl/auto/SurveyCreatorPipeline.py +21 -21
  24. edsl/auto/utilities.py +224 -224
  25. edsl/base/Base.py +279 -289
  26. edsl/config.py +157 -149
  27. edsl/conversation/Conversation.py +290 -290
  28. edsl/conversation/car_buying.py +58 -58
  29. edsl/conversation/chips.py +95 -95
  30. edsl/conversation/mug_negotiation.py +81 -81
  31. edsl/conversation/next_speaker_utilities.py +93 -93
  32. edsl/coop/PriceFetcher.py +54 -54
  33. edsl/coop/__init__.py +2 -2
  34. edsl/coop/coop.py +1028 -958
  35. edsl/coop/utils.py +131 -131
  36. edsl/data/Cache.py +555 -527
  37. edsl/data/CacheEntry.py +233 -228
  38. edsl/data/CacheHandler.py +149 -149
  39. edsl/data/RemoteCacheSync.py +78 -97
  40. edsl/data/SQLiteDict.py +292 -292
  41. edsl/data/__init__.py +4 -4
  42. edsl/data/orm.py +10 -10
  43. edsl/data_transfer_models.py +73 -73
  44. edsl/enums.py +175 -173
  45. edsl/exceptions/BaseException.py +21 -21
  46. edsl/exceptions/__init__.py +54 -54
  47. edsl/exceptions/agents.py +42 -38
  48. edsl/exceptions/cache.py +5 -0
  49. edsl/exceptions/configuration.py +16 -16
  50. edsl/exceptions/coop.py +10 -10
  51. edsl/exceptions/data.py +14 -14
  52. edsl/exceptions/general.py +34 -34
  53. edsl/exceptions/jobs.py +33 -33
  54. edsl/exceptions/language_models.py +63 -63
  55. edsl/exceptions/prompts.py +15 -15
  56. edsl/exceptions/questions.py +91 -91
  57. edsl/exceptions/results.py +29 -29
  58. edsl/exceptions/scenarios.py +22 -22
  59. edsl/exceptions/surveys.py +37 -37
  60. edsl/inference_services/AnthropicService.py +87 -87
  61. edsl/inference_services/AwsBedrock.py +120 -120
  62. edsl/inference_services/AzureAI.py +217 -217
  63. edsl/inference_services/DeepInfraService.py +18 -18
  64. edsl/inference_services/GoogleService.py +148 -156
  65. edsl/inference_services/GroqService.py +20 -20
  66. edsl/inference_services/InferenceServiceABC.py +147 -147
  67. edsl/inference_services/InferenceServicesCollection.py +97 -97
  68. edsl/inference_services/MistralAIService.py +123 -123
  69. edsl/inference_services/OllamaService.py +18 -18
  70. edsl/inference_services/OpenAIService.py +224 -224
  71. edsl/inference_services/PerplexityService.py +163 -0
  72. edsl/inference_services/TestService.py +89 -89
  73. edsl/inference_services/TogetherAIService.py +170 -170
  74. edsl/inference_services/models_available_cache.py +118 -118
  75. edsl/inference_services/rate_limits_cache.py +25 -25
  76. edsl/inference_services/registry.py +41 -39
  77. edsl/inference_services/write_available.py +10 -10
  78. edsl/jobs/Answers.py +56 -56
  79. edsl/jobs/Jobs.py +898 -1347
  80. edsl/jobs/JobsChecks.py +147 -0
  81. edsl/jobs/JobsPrompts.py +268 -0
  82. edsl/jobs/JobsRemoteInferenceHandler.py +239 -0
  83. edsl/jobs/__init__.py +1 -1
  84. edsl/jobs/buckets/BucketCollection.py +63 -63
  85. edsl/jobs/buckets/ModelBuckets.py +65 -65
  86. edsl/jobs/buckets/TokenBucket.py +251 -248
  87. edsl/jobs/interviews/Interview.py +661 -661
  88. edsl/jobs/interviews/InterviewExceptionCollection.py +99 -99
  89. edsl/jobs/interviews/InterviewExceptionEntry.py +186 -186
  90. edsl/jobs/interviews/InterviewStatistic.py +63 -63
  91. edsl/jobs/interviews/InterviewStatisticsCollection.py +25 -25
  92. edsl/jobs/interviews/InterviewStatusDictionary.py +78 -78
  93. edsl/jobs/interviews/InterviewStatusLog.py +92 -92
  94. edsl/jobs/interviews/ReportErrors.py +66 -66
  95. edsl/jobs/interviews/interview_status_enum.py +9 -9
  96. edsl/jobs/runners/JobsRunnerAsyncio.py +466 -338
  97. edsl/jobs/runners/JobsRunnerStatus.py +330 -332
  98. edsl/jobs/tasks/QuestionTaskCreator.py +242 -242
  99. edsl/jobs/tasks/TaskCreators.py +64 -64
  100. edsl/jobs/tasks/TaskHistory.py +450 -442
  101. edsl/jobs/tasks/TaskStatusLog.py +23 -23
  102. edsl/jobs/tasks/task_status_enum.py +163 -163
  103. edsl/jobs/tokens/InterviewTokenUsage.py +27 -27
  104. edsl/jobs/tokens/TokenUsage.py +34 -34
  105. edsl/language_models/KeyLookup.py +30 -30
  106. edsl/language_models/LanguageModel.py +668 -706
  107. edsl/language_models/ModelList.py +155 -102
  108. edsl/language_models/RegisterLanguageModelsMeta.py +184 -184
  109. edsl/language_models/__init__.py +3 -3
  110. edsl/language_models/fake_openai_call.py +15 -15
  111. edsl/language_models/fake_openai_service.py +61 -61
  112. edsl/language_models/registry.py +190 -137
  113. edsl/language_models/repair.py +156 -156
  114. edsl/language_models/unused/ReplicateBase.py +83 -83
  115. edsl/language_models/utilities.py +64 -64
  116. edsl/notebooks/Notebook.py +258 -259
  117. edsl/notebooks/__init__.py +1 -1
  118. edsl/prompts/Prompt.py +362 -357
  119. edsl/prompts/__init__.py +2 -2
  120. edsl/questions/AnswerValidatorMixin.py +289 -289
  121. edsl/questions/QuestionBase.py +664 -656
  122. edsl/questions/QuestionBaseGenMixin.py +161 -161
  123. edsl/questions/QuestionBasePromptsMixin.py +217 -234
  124. edsl/questions/QuestionBudget.py +227 -227
  125. edsl/questions/QuestionCheckBox.py +359 -359
  126. edsl/questions/QuestionExtract.py +182 -183
  127. edsl/questions/QuestionFreeText.py +114 -114
  128. edsl/questions/QuestionFunctional.py +166 -159
  129. edsl/questions/QuestionList.py +231 -231
  130. edsl/questions/QuestionMultipleChoice.py +286 -286
  131. edsl/questions/QuestionNumerical.py +153 -153
  132. edsl/questions/QuestionRank.py +324 -324
  133. edsl/questions/Quick.py +41 -41
  134. edsl/questions/RegisterQuestionsMeta.py +71 -71
  135. edsl/questions/ResponseValidatorABC.py +174 -174
  136. edsl/questions/SimpleAskMixin.py +73 -73
  137. edsl/questions/__init__.py +26 -26
  138. edsl/questions/compose_questions.py +98 -98
  139. edsl/questions/decorators.py +21 -21
  140. edsl/questions/derived/QuestionLikertFive.py +76 -76
  141. edsl/questions/derived/QuestionLinearScale.py +87 -87
  142. edsl/questions/derived/QuestionTopK.py +93 -91
  143. edsl/questions/derived/QuestionYesNo.py +82 -82
  144. edsl/questions/descriptors.py +413 -413
  145. edsl/questions/prompt_templates/question_budget.jinja +13 -13
  146. edsl/questions/prompt_templates/question_checkbox.jinja +32 -32
  147. edsl/questions/prompt_templates/question_extract.jinja +11 -11
  148. edsl/questions/prompt_templates/question_free_text.jinja +3 -3
  149. edsl/questions/prompt_templates/question_linear_scale.jinja +11 -11
  150. edsl/questions/prompt_templates/question_list.jinja +17 -17
  151. edsl/questions/prompt_templates/question_multiple_choice.jinja +33 -33
  152. edsl/questions/prompt_templates/question_numerical.jinja +36 -36
  153. edsl/questions/question_registry.py +177 -147
  154. edsl/questions/settings.py +12 -12
  155. edsl/questions/templates/budget/answering_instructions.jinja +7 -7
  156. edsl/questions/templates/budget/question_presentation.jinja +7 -7
  157. edsl/questions/templates/checkbox/answering_instructions.jinja +10 -10
  158. edsl/questions/templates/checkbox/question_presentation.jinja +22 -22
  159. edsl/questions/templates/extract/answering_instructions.jinja +7 -7
  160. edsl/questions/templates/likert_five/answering_instructions.jinja +10 -10
  161. edsl/questions/templates/likert_five/question_presentation.jinja +11 -11
  162. edsl/questions/templates/linear_scale/answering_instructions.jinja +5 -5
  163. edsl/questions/templates/linear_scale/question_presentation.jinja +5 -5
  164. edsl/questions/templates/list/answering_instructions.jinja +3 -3
  165. edsl/questions/templates/list/question_presentation.jinja +5 -5
  166. edsl/questions/templates/multiple_choice/answering_instructions.jinja +9 -9
  167. edsl/questions/templates/multiple_choice/question_presentation.jinja +11 -11
  168. edsl/questions/templates/numerical/answering_instructions.jinja +6 -6
  169. edsl/questions/templates/numerical/question_presentation.jinja +6 -6
  170. edsl/questions/templates/rank/answering_instructions.jinja +11 -11
  171. edsl/questions/templates/rank/question_presentation.jinja +15 -15
  172. edsl/questions/templates/top_k/answering_instructions.jinja +8 -8
  173. edsl/questions/templates/top_k/question_presentation.jinja +22 -22
  174. edsl/questions/templates/yes_no/answering_instructions.jinja +6 -6
  175. edsl/questions/templates/yes_no/question_presentation.jinja +11 -11
  176. edsl/results/CSSParameterizer.py +108 -0
  177. edsl/results/Dataset.py +424 -293
  178. edsl/results/DatasetExportMixin.py +731 -717
  179. edsl/results/DatasetTree.py +275 -145
  180. edsl/results/Result.py +465 -450
  181. edsl/results/Results.py +1165 -1071
  182. edsl/results/ResultsDBMixin.py +238 -238
  183. edsl/results/ResultsExportMixin.py +43 -43
  184. edsl/results/ResultsFetchMixin.py +33 -33
  185. edsl/results/ResultsGGMixin.py +121 -121
  186. edsl/results/ResultsToolsMixin.py +98 -98
  187. edsl/results/Selector.py +135 -135
  188. edsl/results/TableDisplay.py +198 -0
  189. edsl/results/__init__.py +2 -2
  190. edsl/results/table_display.css +78 -0
  191. edsl/results/tree_explore.py +115 -115
  192. edsl/scenarios/FileStore.py +632 -458
  193. edsl/scenarios/Scenario.py +601 -546
  194. edsl/scenarios/ScenarioHtmlMixin.py +64 -64
  195. edsl/scenarios/ScenarioJoin.py +127 -0
  196. edsl/scenarios/ScenarioList.py +1287 -1112
  197. edsl/scenarios/ScenarioListExportMixin.py +52 -52
  198. edsl/scenarios/ScenarioListPdfMixin.py +261 -261
  199. edsl/scenarios/__init__.py +4 -4
  200. edsl/shared.py +1 -1
  201. edsl/study/ObjectEntry.py +173 -173
  202. edsl/study/ProofOfWork.py +113 -113
  203. edsl/study/SnapShot.py +80 -80
  204. edsl/study/Study.py +528 -528
  205. edsl/study/__init__.py +4 -4
  206. edsl/surveys/DAG.py +148 -148
  207. edsl/surveys/Memory.py +31 -31
  208. edsl/surveys/MemoryPlan.py +244 -244
  209. edsl/surveys/Rule.py +326 -330
  210. edsl/surveys/RuleCollection.py +387 -387
  211. edsl/surveys/Survey.py +1801 -1795
  212. edsl/surveys/SurveyCSS.py +261 -261
  213. edsl/surveys/SurveyExportMixin.py +259 -259
  214. edsl/surveys/SurveyFlowVisualizationMixin.py +179 -121
  215. edsl/surveys/SurveyQualtricsImport.py +284 -284
  216. edsl/surveys/__init__.py +3 -3
  217. edsl/surveys/base.py +53 -53
  218. edsl/surveys/descriptors.py +56 -56
  219. edsl/surveys/instructions/ChangeInstruction.py +49 -47
  220. edsl/surveys/instructions/Instruction.py +65 -51
  221. edsl/surveys/instructions/InstructionCollection.py +77 -77
  222. edsl/templates/error_reporting/base.html +23 -23
  223. edsl/templates/error_reporting/exceptions_by_model.html +34 -34
  224. edsl/templates/error_reporting/exceptions_by_question_name.html +16 -16
  225. edsl/templates/error_reporting/exceptions_by_type.html +16 -16
  226. edsl/templates/error_reporting/interview_details.html +115 -115
  227. edsl/templates/error_reporting/interviews.html +19 -10
  228. edsl/templates/error_reporting/overview.html +4 -4
  229. edsl/templates/error_reporting/performance_plot.html +1 -1
  230. edsl/templates/error_reporting/report.css +73 -73
  231. edsl/templates/error_reporting/report.html +117 -117
  232. edsl/templates/error_reporting/report.js +25 -25
  233. edsl/tools/__init__.py +1 -1
  234. edsl/tools/clusters.py +192 -192
  235. edsl/tools/embeddings.py +27 -27
  236. edsl/tools/embeddings_plotting.py +118 -118
  237. edsl/tools/plotting.py +112 -112
  238. edsl/tools/summarize.py +18 -18
  239. edsl/utilities/SystemInfo.py +28 -28
  240. edsl/utilities/__init__.py +22 -22
  241. edsl/utilities/ast_utilities.py +25 -25
  242. edsl/utilities/data/Registry.py +6 -6
  243. edsl/utilities/data/__init__.py +1 -1
  244. edsl/utilities/data/scooter_results.json +1 -1
  245. edsl/utilities/decorators.py +77 -77
  246. edsl/utilities/gcp_bucket/cloud_storage.py +96 -96
  247. edsl/utilities/interface.py +627 -627
  248. edsl/{conjure → utilities}/naming_utilities.py +263 -263
  249. edsl/utilities/repair_functions.py +28 -28
  250. edsl/utilities/restricted_python.py +70 -70
  251. edsl/utilities/utilities.py +424 -409
  252. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/LICENSE +21 -21
  253. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/METADATA +2 -1
  254. edsl-0.1.38.dist-info/RECORD +277 -0
  255. edsl/conjure/AgentConstructionMixin.py +0 -160
  256. edsl/conjure/Conjure.py +0 -62
  257. edsl/conjure/InputData.py +0 -659
  258. edsl/conjure/InputDataCSV.py +0 -48
  259. edsl/conjure/InputDataMixinQuestionStats.py +0 -182
  260. edsl/conjure/InputDataPyRead.py +0 -91
  261. edsl/conjure/InputDataSPSS.py +0 -8
  262. edsl/conjure/InputDataStata.py +0 -8
  263. edsl/conjure/QuestionOptionMixin.py +0 -76
  264. edsl/conjure/QuestionTypeMixin.py +0 -23
  265. edsl/conjure/RawQuestion.py +0 -65
  266. edsl/conjure/SurveyResponses.py +0 -7
  267. edsl/conjure/__init__.py +0 -9
  268. edsl/conjure/examples/placeholder.txt +0 -0
  269. edsl/conjure/utilities.py +0 -201
  270. edsl-0.1.37.dev6.dist-info/RECORD +0 -283
  271. {edsl-0.1.37.dev6.dist-info → edsl-0.1.38.dist-info}/WHEEL +0 -0
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