Initial llama3 node

2024-06-23 20:27:41 +00:00
parent 082a18f46c
commit 868038477f
1 changed files with 90 additions and 0 deletions
@@ -0,0 +1,90 @@
+import replicate
+
+
+class Llama3Replicate:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "prompt": (
+                    "STRING",
+                    {"default": "", "multiline": True, "dynamicPrompts": True},
+                ),
+                "system_prompt": (
+                    "STRING",
+                    {
+                        "default": "You are a helpful assistant",
+                        "multiline": True,
+                        "dynamicPrompts": True,
+                    },
+                ),
+                "top_p": (
+                    "FLOAT",
+                    {"default": 0.95, "max": 1.0, "min": -1.0},
+                ),
+                "top_k": (
+                    "INT",
+                    {"default": 0, "min": -1},
+                ),
+                "max_tokens": (
+                    "INT",
+                    {"default": 512, "min": 1},
+                ),
+                "min_tokens": (
+                    "INT",
+                    {"default": 0, "min": 0},
+                ),
+                "temperature": (
+                    "FLOAT",
+                    {"default": 0.7, "max": 5.0, "min": 0.0},
+                ),
+                "length_penalty": (
+                    "FLOAT",
+                    {"default": 1.0, "max": 5.0, "min": 0.0},
+                ),
+                "presence_penalty": (
+                    "FLOAT",
+                    {"default": 0},
+                ),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "run_llama3_replicate"
+    CATEGORY = "Replicate"
+
+    def run_llama3_replicate(
+        self,
+        top_p,
+        top_k,
+        prompt,
+        system_prompt,
+        max_tokens,
+        min_tokens,
+        temperature,
+        length_penalty,
+        presence_penalty,
+        seed,
+    ):
+        input = {
+            "system_prompt": system_prompt,
+            "prompt": prompt,
+            "top_p": top_p,
+            "top_k": top_k,
+            "max_tokens": max_tokens,
+            "min_tokens": min_tokens,
+            "temperature": temperature,
+            "length_penalty": length_penalty,
+            "presence_penalty": presence_penalty,
+            "seed": seed,
+        }
+
+        output = replicate.run("meta/meta-llama-3-70b-instruct", input=input)
+        output = "".join(output).strip()
+        return (output,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Llama 3 Replicate": Llama3Replicate,
+}