diff --git a/app/(protected)/project/[project_id]/annotations/annotations.tsx b/app/(protected)/project/[project_id]/annotations/annotations.tsx
index 9716a3bf..d89445dd 100644
--- a/app/(protected)/project/[project_id]/annotations/annotations.tsx
+++ b/app/(protected)/project/[project_id]/annotations/annotations.tsx
@@ -181,15 +181,17 @@ export default function Annotations({ email }: { email: string }) {
           const spanId = row.original.span_id;
           const prompts = row.getValue("input") as string[];
           const responses = row.getValue("output") as string[];
+          const model = row.getValue("model") as string;
 
           let input = "";
           if (prompts && prompts.length > 0) {
             // get last item in prompts
             const lastItem = prompts[prompts.length - 1];
+            const parsedLastItem = JSON.parse(lastItem);
             input =
-              JSON.parse(lastItem)[0]?.content ||
-              JSON.parse(lastItem)[0]?.message?.content ||
-              JSON.parse(lastItem)[0]?.text;
+              parsedLastItem[parsedLastItem.length - 1]?.content ||
+              parsedLastItem[parsedLastItem.length - 1]?.message?.content ||
+              parsedLastItem[parsedLastItem.length - 1]?.text;
 
             // check if input is not a string
             if (typeof input !== "string") {
@@ -201,10 +203,11 @@ export default function Annotations({ email }: { email: string }) {
           if (responses && responses.length > 0) {
             // get last item in responses
             const lastItem = responses[responses.length - 1];
+            const parsedLastItem = JSON.parse(lastItem);
             output =
-              JSON.parse(lastItem)[0]?.message?.content ||
-              JSON.parse(lastItem)[0]?.text ||
-              JSON.parse(lastItem)[0]?.content;
+              parsedLastItem[parsedLastItem.length - 1]?.message?.content ||
+              parsedLastItem[parsedLastItem.length - 1]?.text ||
+              parsedLastItem[parsedLastItem.length - 1]?.content;
 
             // check if output is not a string
             if (typeof output !== "string") {
@@ -221,9 +224,10 @@ export default function Annotations({ email }: { email: string }) {
                   spanId,
                   input,
                   output,
+                  model,
                 };
                 if (checked) {
-                  setSelectedData([...selectedData, checkedData]);
+                  setSelectedData((prev) => [...prev, checkedData]);
                 } else {
                   setSelectedData(
                     selectedData.filter((d) => d.spanId !== spanId)
diff --git a/app/(protected)/project/[project_id]/datasets/dataset/[dataset_id]/page.tsx b/app/(protected)/project/[project_id]/datasets/dataset/[dataset_id]/page.tsx
index a34a88cd..7112eaf5 100644
--- a/app/(protected)/project/[project_id]/datasets/dataset/[dataset_id]/page.tsx
+++ b/app/(protected)/project/[project_id]/datasets/dataset/[dataset_id]/page.tsx
@@ -1,22 +1,51 @@
 "use client";
 
+import { ExpandingTextArea } from "@/components/playground/common";
 import { CreateData } from "@/components/project/dataset/create-data";
 import DatasetRowSkeleton from "@/components/project/dataset/dataset-row-skeleton";
-import { EditData } from "@/components/project/dataset/edit-data";
+import { DeleteData } from "@/components/project/dataset/delete-data";
 import { DownloadDataset } from "@/components/shared/download-dataset";
-import { Spinner } from "@/components/shared/spinner";
+import RowSkeleton from "@/components/shared/row-skeleton";
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuCheckboxItem,
+  DropdownMenuContent,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
 import { Separator } from "@/components/ui/separator";
+import {
+  Table,
+  TableBody,
+  TableCell,
+  TableHead,
+  TableHeader,
+  TableRow,
+} from "@/components/ui/table";
 import { EVALUATIONS_DOCS_URL, PAGE_SIZE } from "@/lib/constants";
+import { cn } from "@/lib/utils";
 import { Data } from "@prisma/client";
-import { ArrowTopRightIcon } from "@radix-ui/react-icons";
-import { ChevronLeft, FlaskConical } from "lucide-react";
+import { ArrowTopRightIcon, ResetIcon } from "@radix-ui/react-icons";
+import {
+  ColumnDef,
+  flexRender,
+  getCoreRowModel,
+  useReactTable,
+  VisibilityState,
+} from "@tanstack/react-table";
+import {
+  ChevronDown,
+  ChevronLeft,
+  FlaskConical,
+  MoveDiagonal,
+  X,
+} from "lucide-react";
 import Link from "next/link";
 import { useParams } from "next/navigation";
-import { useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
 import { useBottomScrollListener } from "react-bottom-scroll-listener";
-import { useQuery } from "react-query";
+import { useQuery, useQueryClient } from "react-query";
 import { toast } from "sonner";
 
 export default function Dataset() {
@@ -24,18 +53,42 @@ export default function Dataset() {
   const dataset_id = useParams()?.dataset_id as string;
   const [page, setPage] = useState<number>(1);
   const [totalPages, setTotalPages] = useState<number>(1);
-  const [showLoader, setShowLoader] = useState(false);
+  const [showBottomLoader, setShowBottomLoader] = useState(false);
   const [currentData, setCurrentData] = useState<Data[]>([]);
+  const [tableState, setTableState] = useState<any>({
+    pagination: {
+      pageIndex: 0,
+      pageSize: 100,
+    },
+  });
+  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({});
+  const [openDropdown, setOpenDropdown] = useState(false);
 
-  useBottomScrollListener(() => {
-    if (fetchDataset.isRefetching) {
-      return;
-    }
-    if (page <= totalPages) {
-      setShowLoader(true);
-      fetchDataset.refetch();
+  useEffect(() => {
+    if (typeof window !== "undefined") {
+      try {
+        const storedState = window.localStorage.getItem(
+          "preferences.dataset.table-view"
+        );
+        if (storedState) {
+          const parsedState = JSON.parse(storedState);
+          setTableState((prevState: any) => ({
+            ...prevState,
+            ...parsedState,
+            pagination: {
+              ...prevState.pagination,
+              ...parsedState.pagination,
+            },
+          }));
+          if (parsedState.columnVisibility) {
+            setColumnVisibility(parsedState.columnVisibility);
+          }
+        }
+      } catch (error) {
+        console.error("Error parsing stored table state:", error);
+      }
     }
-  });
+  }, []);
 
   const fetchDataset = useQuery({
     queryKey: [dataset_id],
@@ -63,29 +116,269 @@ export default function Dataset() {
 
       // Merge the new data with the existing data
       if (currentData.length > 0) {
-        const updatedData = [...currentData, ...newData];
+        const uniqueData = [
+          ...currentData.filter(
+            (item) => !newData.some((existing) => existing.id === item.id)
+          ),
+          ...newData,
+        ];
 
-        // Remove duplicates
-        const uniqueData = updatedData.filter(
-          (v: any, i: number, a: any) =>
-            a.findIndex((t: any) => t.id === v.id) === i
+        // sort new data by created_at
+        uniqueData.sort(
+          (a: Data, b: Data) =>
+            new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()
         );
 
         setCurrentData(uniqueData);
       } else {
+        // sort new data by created_at
+        newData.sort(
+          (a: Data, b: Data) =>
+            new Date(b.createdAt).getTime() - new Date(a.createdAt).getTime()
+        );
         setCurrentData(newData);
       }
-      setShowLoader(false);
+
+      setShowBottomLoader(false);
     },
     onError: (error) => {
-      setShowLoader(false);
+      setShowBottomLoader(false);
       toast.error("Failed to fetch dataset", {
         description: error instanceof Error ? error.message : String(error),
       });
     },
   });
 
-  if (fetchDataset.isLoading || !fetchDataset.data || !currentData) {
+  const scrollableDivRef = useBottomScrollListener(() => {
+    if (fetchDataset.isRefetching) {
+      return;
+    }
+    if (page <= totalPages) {
+      setShowBottomLoader(true);
+      fetchDataset.refetch();
+    }
+  });
+
+  const DataComponent = ({
+    label,
+    value,
+    id,
+    editable = false,
+  }: {
+    label: string;
+    value: string;
+    id?: string;
+    editable?: boolean;
+  }) => {
+    const [expandedView, setExpandedView] = useState(false);
+    const [editMode, setEditMode] = useState(false);
+    const [changedValue, setChangedValue] = useState(value);
+    const saveButtonRef = useRef<HTMLButtonElement>(null);
+    const queryClient = useQueryClient();
+    return (
+      <div className="flex flex-col gap-2">
+        {!editMode && !expandedView && value && (
+          <MoveDiagonal
+            className="self-end h-4 w-4 hover:bg-primary-foreground hover:text-primary cursor-pointer text-muted-foreground"
+            onClick={(e) => {
+              e.stopPropagation();
+              setExpandedView(!expandedView);
+            }}
+          />
+        )}
+        {!editMode && expandedView && value && (
+          <X
+            className="self-end h-4 w-4 hover:bg-primary-foreground hover:text-primary cursor-pointer text-muted-foreground"
+            onClick={(e) => {
+              e.stopPropagation();
+              setExpandedView(!expandedView);
+            }}
+          />
+        )}
+        {editable && editMode && (
+          <ExpandingTextArea
+            onChange={(value: string) => {
+              setChangedValue(value);
+            }}
+            value={changedValue}
+            setFocusing={setEditMode}
+            saveButtonRef={saveButtonRef}
+            saveButtonLabel="Save"
+            handleSave={async () => {
+              try {
+                await fetch("/api/data", {
+                  method: "PUT",
+                  headers: {
+                    "Content-Type": "application/json",
+                  },
+                  body: JSON.stringify({
+                    id: id,
+                    [label]: changedValue,
+                  }),
+                });
+                await queryClient.invalidateQueries({ queryKey: [dataset_id] });
+                toast("Data saved!", {
+                  description: "Your data has been saved.",
+                });
+              } catch (error: any) {
+                toast("Error saving your dataset!", {
+                  description: `There was an error saving your dataset: ${error.message}`,
+                });
+              } finally {
+                setEditMode(false);
+              }
+            }}
+          />
+        )}
+        {(!editable || !editMode) && (
+          <p
+            onClick={() => editable && setEditMode(true)}
+            className={cn(
+              expandedView ? "" : "h-20",
+              "text-sm overflow-y-scroll"
+            )}
+          >
+            {value}
+          </p>
+        )}
+      </div>
+    );
+  };
+
+  const columns: ColumnDef<Data>[] = [
+    {
+      accessorKey: "id",
+      header: "ID",
+      cell: ({ row }) => {
+        return (
+          <p className="overflow-x-scroll text-xs text-muted-foreground">
+            {row.getValue("id")}
+          </p>
+        );
+      },
+    },
+    {
+      size: 500,
+      accessorKey: "input",
+      enableResizing: true,
+      header: "Input",
+      cell: ({ row }) => {
+        const input = row.getValue("input") as string;
+        return <DataComponent value={input} label={"input"} />;
+      },
+    },
+    {
+      size: 500,
+      accessorKey: "output",
+      enableResizing: true,
+      header: "Output",
+      cell: ({ row }) => {
+        const output = row.getValue("output") as string;
+        return <DataComponent value={output} label={"output"} />;
+      },
+    },
+    {
+      size: 500,
+      accessorKey: "expectedOutput",
+      enableResizing: true,
+      header: "Expected Output",
+      cell: ({ row }) => {
+        const expectedOutput = row.getValue("expectedOutput") as string;
+        const id = row.getValue("id") as string;
+        return (
+          <DataComponent
+            editable={true}
+            value={expectedOutput}
+            id={id}
+            label={"expectedOutput"}
+          />
+        );
+      },
+    },
+    {
+      size: 100,
+      accessorKey: "model",
+      enableResizing: true,
+      header: "Model",
+      cell: ({ row }) => {
+        const model = row.getValue("model") as string;
+        return <p className="text-sm h-10 overflow-y-scroll">{model}</p>;
+      },
+    },
+    {
+      size: 300,
+      accessorKey: "note",
+      enableResizing: true,
+      header: "Note",
+      cell: ({ row }) => {
+        const note = row.getValue("note") as string;
+        const id = row.getValue("id") as string;
+        return (
+          <DataComponent editable={true} value={note} id={id} label={"note"} />
+        );
+      },
+    },
+    {
+      size: 50,
+      accessorKey: "delete",
+      enableResizing: true,
+      header: "Delete",
+      cell: ({ row }) => {
+        return <DeleteData id={row.getValue("id") as string} />;
+      },
+    },
+  ];
+
+  const table = useReactTable({
+    data: currentData,
+    columns,
+    getCoreRowModel: getCoreRowModel(),
+    initialState: {
+      ...tableState,
+      pagination: tableState.pagination,
+      columnVisibility,
+    },
+    state: {
+      ...tableState,
+      pagination: tableState.pagination,
+      columnVisibility,
+    },
+    onStateChange: (newState: any) => {
+      setTableState((prevState: any) => ({
+        ...newState,
+        pagination: newState.pagination || prevState.pagination,
+      }));
+      const currState = table.getState();
+      localStorage.setItem(
+        "preferences.dataset.table-view",
+        JSON.stringify(currState)
+      );
+    },
+    onColumnVisibilityChange: (newVisibility) => {
+      setColumnVisibility(newVisibility);
+      const currState = table.getState();
+      localStorage.setItem(
+        "preferences.dataset.table-view",
+        JSON.stringify(currState)
+      );
+    },
+    enableColumnResizing: true,
+    columnResizeMode: "onChange",
+    manualPagination: true, // Add this if you're handling pagination yourself
+  });
+
+  const columnSizeVars = useMemo(() => {
+    const headers = table.getFlatHeaders();
+    const colSizes: { [key: string]: number } = {};
+    for (let i = 0; i < headers.length; i++) {
+      const header = headers[i]!;
+      colSizes[`--header-${header.id}-size`] = header.getSize();
+      colSizes[`--col-${header.column.id}-size`] = header.column.getSize();
+    }
+    return colSizes;
+  }, [table.getState().columnSizingInfo, table.getState().columnSizing]);
+
+  if (fetchDataset.isLoading || !fetchDataset.data) {
     return <PageSkeleton />;
   } else {
     return (
@@ -114,51 +407,123 @@ export default function Dataset() {
                 <ArrowTopRightIcon className="ml-1 h-4 w-4" />
               </Button>
             </Link>
+            <DropdownMenu open={openDropdown} onOpenChange={setOpenDropdown}>
+              <DropdownMenuTrigger asChild>
+                <Button variant="outline" className="ml-auto">
+                  Columns
+                  <ChevronDown size={16} className="ml-2" />
+                </Button>
+              </DropdownMenuTrigger>
+              <DropdownMenuContent className="h-40 overflow-y-visible">
+                {table
+                  .getAllColumns()
+                  .filter((column) => column.getCanHide())
+                  .map((column) => {
+                    return (
+                      <DropdownMenuCheckboxItem
+                        key={column.id}
+                        className="capitalize"
+                        checked={column.getIsVisible()}
+                        onSelect={() => setOpenDropdown(true)}
+                        onCheckedChange={(value) =>
+                          column.toggleVisibility(!!value)
+                        }
+                      >
+                        {column.columnDef.header?.toString()}
+                      </DropdownMenuCheckboxItem>
+                    );
+                  })}
+              </DropdownMenuContent>
+            </DropdownMenu>
+            <Button
+              size={"icon"}
+              variant={"destructive"}
+              onClick={() => {
+                setColumnVisibility({});
+                setTableState({});
+                localStorage.removeItem("preferences.dataset.table-view");
+              }}
+            >
+              <ResetIcon className="w-4 h-4" />
+            </Button>
           </div>
         </div>
-        <div className="flex flex-col gap-3 rounded-md border border-muted max-h-screen overflow-y-scroll">
-          <div className="grid grid-cols-5 items-center justify-stretch gap-3 py-3 px-4 bg-muted">
-            <p className="text-xs font-medium">Created at</p>
-            <p className="text-xs font-medium">Input</p>
-            <p className="text-xs font-medium">Output</p>
-            <p className="text-xs font-medium text-end">Note</p>
-          </div>
-          {fetchDataset.isLoading && currentData?.length === 0 && (
-            <div className="flex items-center justify-center">
+        <div
+          className="rounded-md border flex flex-col relative h-[75vh] overflow-y-scroll"
+          ref={scrollableDivRef as any}
+        >
+          {!fetchDataset.isLoading && currentData.length === 0 && (
+            <div className="flex items-center justify-center mt-24">
               <p className="text-muted-foreground">
                 No data found in this dataset
               </p>
             </div>
           )}
           {!fetchDataset.isLoading &&
-            currentData.length > 0 &&
-            currentData.map((data: any, i: number) => {
-              return (
-                <div className="flex flex-col" key={i}>
-                  <div className="grid grid-cols-5 items-start justify-stretch gap-3 py-3 px-4">
-                    <p className="text-xs">{data.createdAt}</p>
-                    <p className="text-xs h-12 overflow-y-scroll">
-                      {data.input}
-                    </p>
-                    <p className="text-xs h-12 overflow-y-scroll">
-                      {data.output}
-                    </p>
-                    <p className="text-xs text-end">{data.note}</p>
-                    <div className="text-end">
-                      <EditData
-                        key={data.id}
-                        idata={data}
-                        datasetId={dataset_id}
-                      />
-                    </div>
-                  </div>
-                  <Separator orientation="horizontal" />
-                </div>
-              );
-            })}
-          {showLoader && (
-            <div className="flex justify-center py-8">
-              <Spinner className="h-8 w-8 text-center" />
+            currentData &&
+            currentData?.length > 0 && (
+              <Table style={{ ...columnSizeVars, width: table.getTotalSize() }}>
+                <TableHeader className="sticky top-0 bg-secondary">
+                  {table.getHeaderGroups().map((headerGroup) => (
+                    <TableRow key={headerGroup.id}>
+                      {headerGroup.headers.map((header) => (
+                        <TableHead
+                          key={header.id}
+                          style={{
+                            width: `calc(var(--header-${header.id}-size) * 1px)`,
+                            position: "relative",
+                          }}
+                        >
+                          {header.isPlaceholder
+                            ? null
+                            : flexRender(
+                                header.column.columnDef.header,
+                                header.getContext()
+                              )}
+                          <div
+                            onDoubleClick={() => header.column.resetSize()}
+                            onMouseDown={header.getResizeHandler()}
+                            onTouchStart={header.getResizeHandler()}
+                            className={`bg-muted-foreground resizer ${
+                              header.column.getIsResizing() ? "isResizing" : ""
+                            }`}
+                          />
+                        </TableHead>
+                      ))}
+                    </TableRow>
+                  ))}
+                </TableHeader>
+                <TableBody>
+                  {table.getRowModel().rows.map((row) => (
+                    <TableRow
+                      className="cursor-pointer"
+                      key={row.id}
+                      data-state={row.getIsSelected() && "selected"}
+                    >
+                      {row.getVisibleCells().map((cell) => (
+                        <TableCell
+                          key={cell.id}
+                          style={{
+                            width: `calc(var(--col-${cell.column.id}-size) * 1px)`,
+                          }}
+                        >
+                          {flexRender(
+                            cell.column.columnDef.cell,
+                            cell.getContext()
+                          )}
+                        </TableCell>
+                      ))}
+                    </TableRow>
+                  ))}
+                </TableBody>
+              </Table>
+            )}
+          {showBottomLoader && (
+            <div className="flex flex-col gap-3">
+              <Separator />
+              {Array.from({ length: 2 }).map((_, index) => (
+                <RowSkeleton key={index} />
+              ))}
             </div>
           )}
         </div>
@@ -183,12 +548,13 @@ function PageSkeleton() {
       </div>
       <div className="flex flex-col gap-3 rounded-md border border-muted max-h-screen overflow-y-scroll">
         <div className="grid grid-cols-5 items-center justify-stretch gap-3 py-3 px-4 bg-muted">
-          <p className="text-xs font-medium">Created at</p>
           <p className="text-xs font-medium">Input</p>
           <p className="text-xs font-medium">Output</p>
+          <p className="text-xs font-medium">Expected Output</p>
+          <p className="text-xs font-medium">Model</p>
           <p className="text-xs font-medium text-end">Note</p>
         </div>
-        {Array.from({ length: 5 }).map((_, index) => (
+        {Array.from({ length: 6 }).map((_, index) => (
           <DatasetRowSkeleton key={index} />
         ))}
       </div>
diff --git a/app/api/data/route.ts b/app/api/data/route.ts
index 721aba11..6130cb6d 100644
--- a/app/api/data/route.ts
+++ b/app/api/data/route.ts
@@ -88,12 +88,12 @@ export async function POST(req: NextRequest) {
 
     const payload = datas.map((data: Data) => {
       const d: any = {
-        input: data.input,
-        output: data.output,
-        contexts: data.contexts || [],
-        annotatedOutput: data.annotatedOutput || "",
-        note: data.note || "",
         spanId: data.spanId || "",
+        input: data.input || "",
+        output: data.output || "",
+        expectedOutput: data.expectedOutput || "",
+        model: data.model || "",
+        note: data.note || "",
         projectId: projectId || "",
         datasetId: datasetId || "",
       };
@@ -106,6 +106,7 @@ export async function POST(req: NextRequest) {
 
     const result = await prisma.data.createMany({
       data: payload,
+      skipDuplicates: true,
     });
 
     return NextResponse.json({
@@ -129,7 +130,7 @@ export async function PUT(req: NextRequest) {
     }
 
     const data = await req.json();
-    const { id, input, output, annotatedOutput, contexts, note } = data;
+    const { id, input, output, expectedOutput, model, note } = data;
 
     const result = await prisma.data.update({
       where: {
@@ -138,8 +139,8 @@ export async function PUT(req: NextRequest) {
       data: {
         input,
         output,
-        annotatedOutput,
-        contexts,
+        expectedOutput,
+        model,
         note,
       },
     });
diff --git a/app/api/dataset/download/route.ts b/app/api/dataset/download/route.ts
index b26f0968..343b0140 100644
--- a/app/api/dataset/download/route.ts
+++ b/app/api/dataset/download/route.ts
@@ -92,7 +92,9 @@ export async function GET(req: NextRequest) {
       result.push({
         input: d.input,
         target: d.output,
-        annotated_output: d.annotatedOutput,
+        expected_output: d.expectedOutput,
+        model: d.model,
+        note: d.note,
       });
     });
 
diff --git a/app/api/metrics/score/route.ts b/app/api/metrics/score/route.ts
index 0d9f0210..b427df0c 100644
--- a/app/api/metrics/score/route.ts
+++ b/app/api/metrics/score/route.ts
@@ -95,7 +95,8 @@ export async function POST(req: NextRequest) {
           dateScoreMap[date][`${evaluation.Test?.name}(${testId})`][0] +
             evaluation.ltUserScore || 0;
         dateScoreMap[date][`${evaluation.Test?.name}(${testId})`][0] = total;
-        dateScoreMap[date][`${evaluation.Test?.name}(${testId})`][1] += 1;
+        dateScoreMap[date][`${evaluation.Test?.name}(${testId})`][1] +=
+          evaluation.Test?.max || 1;
       });
     }
 
diff --git a/components/annotations/annotations-table.tsx b/components/annotations/annotations-table.tsx
index 1e93d16f..e37a4cde 100644
--- a/components/annotations/annotations-table.tsx
+++ b/components/annotations/annotations-table.tsx
@@ -368,7 +368,6 @@ function EvaluationSheet({
   const {
     isError,
     isLoading,
-    isFetching,
     data: evaluations,
   } = useQuery({
     queryKey: ["fetch-evaluation-query", span.span_id],
diff --git a/components/annotations/chart-tabs.tsx b/components/annotations/chart-tabs.tsx
index 188db007..555c8068 100644
--- a/components/annotations/chart-tabs.tsx
+++ b/components/annotations/chart-tabs.tsx
@@ -89,7 +89,7 @@ export function ChartTabs({
             <Skeleton className="h-24 w-24 rounded-full" />
           </div>
         ) : Object.keys(chartData?.scores).length > 0 ? (
-          <div className="flex flex-row gap-4 flex-wrap">
+          <div className="flex flex-row gap-6 flex-wrap">
             {Object.keys(chartData?.scores).map(
               (testId: string, index: number) => {
                 const score = (chartData?.scores[testId] || 0) * 100;
@@ -110,14 +110,14 @@ export function ChartTabs({
                     >
                       <span
                         className={cn(
-                          "text-lg font-semibold text-primary rounded-full h-16 w-16 flex items-center justify-center",
+                          "text-lg font-semibold text-black rounded-full h-16 w-16 flex items-center justify-center",
                           bgColor
                         )}
                       >
                         {(chartData?.scores[testId] || 0) * 100}%
                       </span>
                     </ProgressCircle>
-                    <span className="text-lg font-semibold capitalize">
+                    <span className="text-lg text-center font-semibold capitalize text-wrap w-24">
                       {testName}
                     </span>
                   </div>
diff --git a/components/playground/common.tsx b/components/playground/common.tsx
index e36f5d15..9946f614 100644
--- a/components/playground/common.tsx
+++ b/components/playground/common.tsx
@@ -9,7 +9,6 @@ import {
   OpenAIRole,
 } from "@/lib/types/playground_types";
 import { cn } from "@/lib/utils";
-import { ArrowTopRightIcon } from "@radix-ui/react-icons";
 import { MinusCircleIcon, PlusIcon } from "lucide-react";
 import Link from "next/link";
 import { useEffect, useRef, useState } from "react";
@@ -41,12 +40,14 @@ export function ExpandingTextArea({
   onChange,
   setFocusing,
   saveButtonRef,
+  saveButtonLabel = "Save to Registry",
   handleSave,
 }: {
   value: string;
   onChange: any;
   setFocusing?: any;
   saveButtonRef: React.RefObject<HTMLButtonElement>;
+  saveButtonLabel?: string;
   handleSave: (open: boolean) => void;
 }) {
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
@@ -86,16 +87,16 @@ export function ExpandingTextArea({
         onChange={handleChange}
         style={{ overflowY: "auto", height: "auto" }}
       />
-      <div className="absolute right-0 top-[-10px] py-2">
+      <div className="absolute right-1 top-0 py-2">
         <Button
           className="text-xs"
           size={"sm"}
           variant="outline"
           onClick={() => handleSave(true)}
+          disabled={value === ""}
           ref={saveButtonRef}
         >
-          Save to Registry
-          <ArrowTopRightIcon className="ml-1 w-4 h-4" />
+          {saveButtonLabel}
         </Button>
       </div>
     </div>
diff --git a/components/project/dataset/create-data.tsx b/components/project/dataset/create-data.tsx
index d1da4358..59cad731 100644
--- a/components/project/dataset/create-data.tsx
+++ b/components/project/dataset/create-data.tsx
@@ -44,7 +44,9 @@ export function CreateData({
   const [busy, setBusy] = useState<boolean>(false);
   const schema = z.object({
     input: z.string().min(2, "Too short").max(200, "Too long"),
-    output: z.string().min(2, "Too short").max(2000, "Too long"),
+    output: z.string().min(2, "Too short").max(2000, "Too long").optional(),
+    expectedOutput: z.string().min(2, "Too short").max(2000, "Too long"),
+    model: z.string().min(2, "Too short").max(20, "Too long").optional(),
     note: z.string().max(25, "Too long").optional(),
   });
   const CreateDataForm = useForm({
@@ -85,7 +87,9 @@ export function CreateData({
                     datas: [
                       {
                         input: data.input,
-                        output: data.output,
+                        output: data.output || "",
+                        expectedOutput: data.expectedOutput,
+                        model: data.model || "",
                         note: data.note || "",
                       },
                     ],
@@ -152,6 +156,49 @@ export function CreateData({
                 </FormItem>
               )}
             />
+            <FormField
+              disabled={busy}
+              control={CreateDataForm.control}
+              name="expectedOutput"
+              render={({ field }) => (
+                <FormItem>
+                  <FormLabel>
+                    Expected Output
+                    <Info
+                      information="Expected response to the input data by the LLM."
+                      className="inline-block ml-2"
+                    />
+                  </FormLabel>
+                  <FormControl>
+                    <Input
+                      placeholder="I'm good, how can I help you?"
+                      {...field}
+                    />
+                  </FormControl>
+                  <FormMessage />
+                </FormItem>
+              )}
+            />
+            <FormField
+              disabled={busy}
+              control={CreateDataForm.control}
+              name="model"
+              render={({ field }) => (
+                <FormItem>
+                  <FormLabel>
+                    Model
+                    <Info
+                      information="Model used to generate the output."
+                      className="inline-block ml-2"
+                    />
+                  </FormLabel>
+                  <FormControl>
+                    <Input placeholder="gpt-4o-mini" {...field} />
+                  </FormControl>
+                  <FormMessage />
+                </FormItem>
+              )}
+            />
             <FormField
               disabled={busy}
               control={CreateDataForm.control}
diff --git a/components/project/dataset/delete-data.tsx b/components/project/dataset/delete-data.tsx
new file mode 100644
index 00000000..afbc465d
--- /dev/null
+++ b/components/project/dataset/delete-data.tsx
@@ -0,0 +1,70 @@
+import { Button } from "@/components/ui/button";
+import {
+  Dialog,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
+import { TrashIcon } from "lucide-react";
+import { useState } from "react";
+import { useQueryClient } from "react-query";
+import { toast } from "sonner";
+
+export function DeleteData({ id }: { id: string }) {
+  const queryClient = useQueryClient();
+  const [open, setOpen] = useState<boolean>(false);
+  const [busy, setBusy] = useState<boolean>(false);
+  return (
+    <Dialog open={open} onOpenChange={setOpen}>
+      <DialogTrigger asChild>
+        <Button variant="ghost" size={"icon"}>
+          <TrashIcon className="w-4 h-4" />
+        </Button>
+      </DialogTrigger>
+      <DialogContent className="sm:max-w-[425px]">
+        <DialogHeader>
+          <DialogTitle>Delete Data</DialogTitle>
+          <DialogDescription>
+            Are you sure you want to delete this data?
+          </DialogDescription>
+        </DialogHeader>
+        <DialogFooter>
+          <Button
+            variant="destructive"
+            onClick={async () => {
+              try {
+                setBusy(true);
+                await fetch("/api/data", {
+                  method: "DELETE",
+                  headers: {
+                    "Content-Type": "application/json",
+                  },
+                  body: JSON.stringify({
+                    id: id,
+                  }),
+                });
+                await queryClient.invalidateQueries({ queryKey: [id] });
+                toast("Dataset deleted!", {
+                  description: "Your data set has been deleted.",
+                });
+                setOpen(false);
+              } catch (error: any) {
+                toast("Error deleting your data set!", {
+                  description: `There was an error deleting your data set: ${error.message}`,
+                });
+              } finally {
+                setBusy(false);
+              }
+            }}
+            disabled={busy}
+          >
+            Delete Dataset
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  );
+}
diff --git a/components/project/dataset/edit-data.tsx b/components/project/dataset/edit-data.tsx
deleted file mode 100644
index 6d9a698b..00000000
--- a/components/project/dataset/edit-data.tsx
+++ /dev/null
@@ -1,286 +0,0 @@
-import { Info } from "@/components/shared/info";
-import { Button } from "@/components/ui/button";
-import {
-  Dialog,
-  DialogContent,
-  DialogDescription,
-  DialogFooter,
-  DialogHeader,
-  DialogTitle,
-  DialogTrigger,
-} from "@/components/ui/dialog";
-import {
-  DropdownMenu,
-  DropdownMenuContent,
-  DropdownMenuItem,
-  DropdownMenuTrigger,
-} from "@/components/ui/dropdown-menu";
-import {
-  Form,
-  FormControl,
-  FormField,
-  FormItem,
-  FormLabel,
-  FormMessage,
-} from "@/components/ui/form";
-import { Input } from "@/components/ui/input";
-import { zodResolver } from "@hookform/resolvers/zod";
-import { Data } from "@prisma/client";
-import { DotsHorizontalIcon } from "@radix-ui/react-icons";
-import { EditIcon, TrashIcon } from "lucide-react";
-import { useState } from "react";
-import { useForm } from "react-hook-form";
-import { useQueryClient } from "react-query";
-import { toast } from "sonner";
-import { z } from "zod";
-
-export function EditData({
-  idata,
-  datasetId,
-  className = "w-full text-left p-0 text-muted-foreground hover:text-primary flex items-center",
-}: {
-  idata: Data;
-  datasetId: string;
-  variant?: any;
-  className?: string;
-}) {
-  const queryClient = useQueryClient();
-  const [open, setOpen] = useState<boolean>(false);
-  const [openEdit, setOpenEdit] = useState<boolean>(false);
-  const [openDelete, setOpenDelete] = useState<boolean>(false);
-  const [busy, setBusy] = useState<boolean>(false);
-  const schema = z.object({
-    input: z.string().min(2, "Too short").max(200, "Too long"),
-    output: z.string().min(2, "Too short").max(2000, "Too long"),
-    note: z.string().max(25, "Too long").optional(),
-  });
-  const EditDataForm = useForm({
-    resolver: zodResolver(schema),
-    defaultValues: {
-      input: idata.input || "",
-      output: idata.output || "",
-      note: idata.note || "",
-    },
-  });
-  return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      <DropdownMenu>
-        <DropdownMenuTrigger asChild>
-          <Button variant="outline" size={"icon"}>
-            <DotsHorizontalIcon />
-          </Button>
-        </DropdownMenuTrigger>
-        <DropdownMenuContent className="w-56">
-          <DropdownMenuItem>
-            <DialogTrigger asChild>
-              <button
-                className={className}
-                onClick={(e) => {
-                  e.preventDefault();
-                  e.stopPropagation();
-
-                  // Reset the form to the data values
-                  // to prevent stale data in the form
-                  EditDataForm.reset({
-                    input: idata.input || "",
-                    output: idata.output || "",
-                    note: idata.note || "",
-                  });
-
-                  setOpenDelete(false);
-                  setOpen(true);
-                  setOpenEdit(true);
-                }}
-              >
-                <EditIcon className="h-4 w-4 mr-2" />
-                Edit Data
-              </button>
-            </DialogTrigger>
-          </DropdownMenuItem>
-          <DropdownMenuItem>
-            <DialogTrigger asChild>
-              <button
-                className={className}
-                onClick={(e) => {
-                  e.preventDefault();
-                  e.stopPropagation();
-                  setOpenDelete(true);
-                  setOpen(true);
-                  setOpenEdit(false);
-                }}
-              >
-                <TrashIcon className="h-4 w-4 mr-2" />
-                Delete Data
-              </button>
-            </DialogTrigger>
-          </DropdownMenuItem>
-        </DropdownMenuContent>
-      </DropdownMenu>
-      {openEdit && (
-        <DialogContent className="sm:max-w-[425px]">
-          <DialogHeader>
-            <DialogTitle>Edit Data</DialogTitle>
-            <DialogDescription>
-              Edit the data by filling out the form below.
-            </DialogDescription>
-          </DialogHeader>
-          <Form {...EditDataForm}>
-            <form
-              onSubmit={EditDataForm.handleSubmit(async (data) => {
-                try {
-                  // stringify data.input, data.output if it's an object
-                  if (typeof data.input === "object") {
-                    data.input = JSON.stringify(data.input);
-                  }
-                  if (typeof data.output === "object") {
-                    data.output = JSON.stringify(data.output);
-                  }
-                  setBusy(true);
-                  await fetch("/api/data", {
-                    method: "PUT",
-                    headers: {
-                      "Content-Type": "application/json",
-                    },
-                    body: JSON.stringify({
-                      id: idata.id,
-                      input: data.input,
-                      output: data.output,
-                      note: data.note,
-                    }),
-                  });
-                  await queryClient.invalidateQueries(datasetId);
-                  toast("Data saved!", {
-                    description: "Your data has been saved.",
-                  });
-                  setOpen(false);
-                } catch (error: any) {
-                  toast("Error saving your dataset!", {
-                    description: `There was an error saving your dataset: ${error.message}`,
-                  });
-                } finally {
-                  setBusy(false);
-                }
-              })}
-              className="flex flex-col gap-4"
-            >
-              <FormField
-                disabled={busy}
-                control={EditDataForm.control}
-                name="input"
-                render={({ field }) => (
-                  <FormItem>
-                    <FormLabel>
-                      Input
-                      <Info
-                        information="The input data. Ex: user input"
-                        className="inline-block ml-2"
-                      />
-                    </FormLabel>
-                    <FormControl>
-                      <Input placeholder="Hi, how are you?" {...field} />
-                    </FormControl>
-                    <FormMessage />
-                  </FormItem>
-                )}
-              />
-              <FormField
-                disabled={busy}
-                control={EditDataForm.control}
-                name="output"
-                render={({ field }) => (
-                  <FormItem>
-                    <FormLabel>
-                      Output
-                      <Info
-                        information="Response to the input data by the LLM."
-                        className="inline-block ml-2"
-                      />
-                    </FormLabel>
-                    <FormControl>
-                      <Input
-                        placeholder="I'm good, how can I help you?"
-                        {...field}
-                      />
-                    </FormControl>
-                    <FormMessage />
-                  </FormItem>
-                )}
-              />
-              <FormField
-                disabled={busy}
-                control={EditDataForm.control}
-                name="note"
-                render={({ field }) => (
-                  <FormItem>
-                    <FormLabel>
-                      Note
-                      <Info
-                        information="Any additional notes for the data."
-                        className="inline-block ml-2"
-                      />
-                    </FormLabel>
-                    <FormControl>
-                      <Input
-                        placeholder="User is asking for help."
-                        {...field}
-                      />
-                    </FormControl>
-                    <FormMessage />
-                  </FormItem>
-                )}
-              />
-              <DialogFooter>
-                <Button type="submit" disabled={busy}>
-                  Save
-                </Button>
-              </DialogFooter>
-            </form>
-          </Form>
-        </DialogContent>
-      )}
-      {openDelete && (
-        <DialogContent className="sm:max-w-[425px]">
-          <DialogHeader>
-            <DialogTitle>Delete Data</DialogTitle>
-            <DialogDescription>
-              Are you sure you want to delete this data?
-            </DialogDescription>
-          </DialogHeader>
-          <DialogFooter>
-            <Button
-              variant="destructive"
-              onClick={async () => {
-                try {
-                  setBusy(true);
-                  await fetch("/api/data", {
-                    method: "DELETE",
-                    headers: {
-                      "Content-Type": "application/json",
-                    },
-                    body: JSON.stringify({
-                      id: idata.id,
-                    }),
-                  });
-                  await queryClient.invalidateQueries(datasetId);
-                  toast("Dataset deleted!", {
-                    description: "Your data set has been deleted.",
-                  });
-                  setOpen(false);
-                } catch (error: any) {
-                  toast("Error deleting your data set!", {
-                    description: `There was an error deleting your data set: ${error.message}`,
-                  });
-                } finally {
-                  setBusy(false);
-                }
-              }}
-              disabled={busy}
-            >
-              Delete Dataset
-            </Button>
-          </DialogFooter>
-        </DialogContent>
-      )}
-    </Dialog>
-  );
-}
diff --git a/components/project/traces/trace-sheet.tsx b/components/project/traces/trace-sheet.tsx
index 3cafbaa3..366a31e4 100644
--- a/components/project/traces/trace-sheet.tsx
+++ b/components/project/traces/trace-sheet.tsx
@@ -83,7 +83,8 @@ export function TraceSheet({
                 <div className="flex gap-2 items-center justify-between w-full">
                   <Button
                     className="w-fit"
-                    variant={"secondary"}
+                    size={"sm"}
+                    variant={"outline"}
                     onClick={() => setSpansView("SPANS")}
                   >
                     <ChevronLeft size={16} className="mr-2" />
@@ -92,7 +93,8 @@ export function TraceSheet({
                   <div className="flex gap-2 items-center">
                     <Button
                       className="w-fit"
-                      variant={"secondary"}
+                      size={"sm"}
+                      variant={"outline"}
                       disabled={spansView === "ATTRIBUTES"}
                       onClick={() => setSpansView("ATTRIBUTES")}
                     >
@@ -101,11 +103,12 @@ export function TraceSheet({
                     </Button>
                     <Button
                       className="w-fit"
-                      variant={"secondary"}
+                      size={"sm"}
+                      variant={"outline"}
                       disabled={spansView === "CONVERSATION"}
                       onClick={() => setSpansView("CONVERSATION")}
                     >
-                      <MessageCircle size={16} className="mr-2" />
+                      <MessageCircle size={16} className="mr-2 fill-primary" />
                       LLM Conversations
                     </Button>
                     {includesLanggraph && (
diff --git a/components/project/traces/traces-table.tsx b/components/project/traces/traces-table.tsx
index 6bb56918..38abeba0 100644
--- a/components/project/traces/traces-table.tsx
+++ b/components/project/traces/traces-table.tsx
@@ -30,7 +30,7 @@ import {
 import { ChevronDown, RefreshCwIcon } from "lucide-react";
 import Link from "next/link";
 import { useEffect, useMemo, useState } from "react";
-import TraceRowSkeleton from "../../shared/row-skeleton";
+import RowSkeleton from "../../shared/row-skeleton";
 import { TableSkeleton } from "./table-skeleton";
 import { TraceSheet } from "./trace-sheet";
 import { TracesDownload } from "./traces-download";
@@ -317,7 +317,7 @@ export function TracesTable<TData, TValue>({
           <div className="flex flex-col gap-3">
             <Separator />
             {Array.from({ length: 2 }).map((_, index) => (
-              <TraceRowSkeleton key={index} />
+              <RowSkeleton key={index} />
             ))}
           </div>
         )}
diff --git a/components/shared/add-to-dataset.tsx b/components/shared/add-to-dataset.tsx
index e4bf4fd8..2bb07072 100644
--- a/components/shared/add-to-dataset.tsx
+++ b/components/shared/add-to-dataset.tsx
@@ -1,8 +1,3 @@
-"use client";
-
-import { Check, ChevronsUpDown, PlusIcon } from "lucide-react";
-import * as React from "react";
-
 import { Button } from "@/components/ui/button";
 import {
   Command,
@@ -11,13 +6,6 @@ import {
   CommandInput,
   CommandItem,
 } from "@/components/ui/command";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
-import { cn } from "@/lib/utils";
-
 import {
   Dialog,
   DialogContent,
@@ -28,7 +16,16 @@ import {
   DialogTrigger,
 } from "@/components/ui/dialog";
 import { Label } from "@/components/ui/label";
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover";
+import { cn } from "@/lib/utils";
+import { Skeleton } from "@mui/material";
 import { Dataset } from "@prisma/client";
+import { Check, ChevronsUpDown, PlusIcon } from "lucide-react";
+import * as React from "react";
 import { useQuery, useQueryClient } from "react-query";
 import { toast } from "sonner";
 import { CreateDataset } from "../project/dataset/create";
@@ -140,7 +137,7 @@ export default function DatasetCombobox({
   });
 
   if (fetchDatasets.isLoading || !fetchDatasets.data) {
-    return <div>Loading...</div>; // this componenet isn't being used, will add updated loading later
+    return <Skeleton width={200} height={50} />;
   } else {
     return (
       <Popover open={open} onOpenChange={setOpen}>
diff --git a/prisma/migrations/20240804000535_modify_data_columns_expectedoutput_model/migration.sql b/prisma/migrations/20240804000535_modify_data_columns_expectedoutput_model/migration.sql
new file mode 100644
index 00000000..96c7de01
--- /dev/null
+++ b/prisma/migrations/20240804000535_modify_data_columns_expectedoutput_model/migration.sql
@@ -0,0 +1,12 @@
+/*
+  Warnings:
+
+  - You are about to drop the column `annotatedOutput` on the `Data` table. All the data in the column will be lost.
+  - You are about to drop the column `contexts` on the `Data` table. All the data in the column will be lost.
+
+*/
+-- AlterTable
+ALTER TABLE "Data" DROP COLUMN "annotatedOutput",
+DROP COLUMN "contexts",
+ADD COLUMN     "expectedOutput" TEXT,
+ADD COLUMN     "model" TEXT;
diff --git a/prisma/schema.prisma b/prisma/schema.prisma
index d28b453a..03a10bf7 100644
--- a/prisma/schema.prisma
+++ b/prisma/schema.prisma
@@ -128,20 +128,20 @@ model Prompt {
 }
 
 model Data {
-  id              String       @id @default(cuid())
-  input           String
-  output          String
-  contexts        String[]
-  annotatedOutput String?
-  note            String?
-  spanId          String?
-  Evaluation      Evaluation[]
-  datasetId       String?
-  Dataset         Dataset?     @relation(fields: [datasetId], references: [id], onDelete: Cascade)
-  projectId       String?
-  Project         Project?     @relation(fields: [projectId], references: [id], onDelete: Cascade)
-  createdAt       DateTime     @default(now())
-  updatedAt       DateTime     @updatedAt
+  id             String       @id @default(cuid())
+  spanId         String?
+  input          String
+  output         String
+  expectedOutput String?
+  model          String?
+  note           String?
+  Evaluation     Evaluation[]
+  datasetId      String?
+  Dataset        Dataset?     @relation(fields: [datasetId], references: [id], onDelete: Cascade)
+  projectId      String?
+  Project        Project?     @relation(fields: [projectId], references: [id], onDelete: Cascade)
+  createdAt      DateTime     @default(now())
+  updatedAt      DateTime     @updatedAt
 }
 
 model Dataset {