fix notebook

ml6team · Feb 6, 2024 · 70dbe30 · 70dbe30
1 parent e8e4bf3
commit 70dbe30
Show file tree

Hide file tree

Showing 4 changed files with 64 additions and 2,660 deletions.
diff --git a/src/components/aggregrate_eval_results.py b/src/components/aggregrate_eval_results.py
@@ -5,14 +5,16 @@
 
 
 @lightweight_component(
+    consumes={
+    "context_precision": pa.float32(),
+    "context_relevancy": pa.float32(),
+    },
     produces={
         "metric": pa.string(),
         "score": pa.float32()
     }
 )
 class AggregateResults(DaskTransformComponent):
-    def __init__(self, consumes: dict, **kwargs):
-        self.consumes = consumes
 
     def transform(self, dataframe: dd.DataFrame) -> dd.DataFrame:
         metrics = list(self.consumes.keys())

diff --git a/src/components/retrieve_from_weaviate.py b/src/components/retrieve_from_weaviate.py
@@ -56,9 +56,8 @@ def retrieve_chunks_from_embeddings(self, vector_query: list):
         )
 
         result = query.do()
-        if "data" in result:
-            result_dict = result["data"]["Get"][self.class_name]
-            return [retrieved_chunk["passage"] for retrieved_chunk in result_dict]
+        result_dict = result["data"]["Get"][self.class_name]
+        return [retrieved_chunk["passage"] for retrieved_chunk in result_dict]
 
     def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
 
@@ -67,12 +66,8 @@ def transform(self, dataframe: pd.DataFrame) -> pd.DataFrame:
                 self.retrieve_chunks_from_embeddings,
             )
 
-        elif "prompt" in dataframe.columns:
-            dataframe["retrieved_chunks"] = dataframe["prompt"].apply(
-                self.retrieve_chunks_from_prompts,
-            )
         else:
-            msg = "Dataframe must contain either an 'embedding' column or a 'prompt' column."
+            msg = "Dataframe must contain an 'embedding' column"
             raise ValueError(
                 msg,
             )