From 055f065da13a1d17095774d96980b071ff55db45 Mon Sep 17 00:00:00 2001
From: Vaibhav Srivastav <vaibhavs10@gmail.com>
Date: Sat, 24 Aug 2024 15:57:20 +0200
Subject: [PATCH] add cp/ edge to supported libs.

---
 packages/tasks/src/model-libraries-snippets.ts | 15 +++++++++++++++
 packages/tasks/src/model-libraries.ts          |  8 +++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts
index 3dbb03920..9adc5c023 100644
--- a/packages/tasks/src/model-libraries-snippets.ts
+++ b/packages/tasks/src/model-libraries-snippets.ts
@@ -170,6 +170,21 @@ export const diffusers = (model: ModelData): string[] => {
 	}
 };
 
+export const edge = (model: ModelData): string[] => [
+	`# pip install --no-binary :all: cartesia-pytorch
+from cartesia_pytorch import ReneLMHeadModel
+from transformers import AutoTokenizer
+
+model = ReneLMHeadModel.from_pretrained("${model.id}").half().cuda()
+tokenizer = AutoTokenizer.from_pretrained("allenai/OLMo-1B-hf")
+in_message = ["Rene Descartes was"]
+inputs = tokenizer(in_message, return_tensors="pt")
+outputs = model.generate(inputs.input_ids.cuda(), max_length=50, top_k=100, top_p=0.99)
+out_message = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+print(out_message)
+)`,
+];
+
 export const edsnlp = (model: ModelData): string[] => {
 	const packageName = nameWithoutNamespace(model.id).replaceAll("-", "_");
 	return [
diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts
index 00f7e81f1..6e77d475d 100644
--- a/packages/tasks/src/model-libraries.ts
+++ b/packages/tasks/src/model-libraries.ts
@@ -128,7 +128,7 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
 		snippets: snippets.bm25s,
 		filter: false,
 		countDownloads: `path:"params.index.json"`,
-	},
+	},	
 	champ: {
 		prettyLabel: "Champ",
 		repoName: "Champ",
@@ -179,6 +179,12 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
 		repoName: "doctr",
 		repoUrl: "https://github.com/mindee/doctr",
 	},
+	edge: {
+		prettyLabel: "Edge",
+		repoName: "Edge",
+		repoUrl: "https://github.com/cartesia-ai/edge",
+		snippets: snippets.edge,		
+	},	
 	edsnlp: {
 		prettyLabel: "EDS-NLP",
 		repoName: "edsnlp",