Set default context size to 8192, adjust with a UI slider

undreamai · Sep 26, 2024 · a6f4dce · a6f4dce
1 parent 8e1253f
commit a6f4dce
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 14 deletions.
diff --git a/Editor/PropertyEditor.cs b/Editor/PropertyEditor.cs
@@ -115,13 +115,32 @@ public void ShowPropertiesOfClass(string title, SerializedObject so, List<Type>
             {
                 Attribute floatAttr = GetPropertyAttribute(prop, typeof(FloatAttribute));
                 Attribute intAttr = GetPropertyAttribute(prop, typeof(IntAttribute));
+                Attribute dynamicRangeAttr = GetPropertyAttribute(prop, typeof(DynamicRangeAttribute));
                 if (floatAttr != null)
                 {
-                    EditorGUILayout.Slider(prop, ((FloatAttribute)floatAttr).Min, ((FloatAttribute)floatAttr).Max, new GUIContent(prop.displayName));
+                    FloatAttribute attr = (FloatAttribute)floatAttr;
+                    EditorGUILayout.Slider(prop, attr.Min, attr.Max, new GUIContent(prop.displayName));
                 }
                 else if (intAttr != null)
                 {
-                    EditorGUILayout.IntSlider(prop, ((IntAttribute)intAttr).Min, ((IntAttribute)intAttr).Max, new GUIContent(prop.displayName));
+                    IntAttribute attr = (IntAttribute)intAttr;
+                    EditorGUILayout.IntSlider(prop, attr.Min, attr.Max, new GUIContent(prop.displayName));
+                }
+                else if (dynamicRangeAttr != null)
+                {
+                    DynamicRangeAttribute attr = (DynamicRangeAttribute)dynamicRangeAttr;
+                    if (attr.intOrFloat)
+                    {
+                        float minValue = so.FindProperty(attr.minVariable).floatValue;
+                        float maxValue = so.FindProperty(attr.maxVariable).floatValue;
+                        EditorGUILayout.Slider(prop, minValue, maxValue, new GUIContent(prop.displayName));
+                    }
+                    else
+                    {
+                        int minValue = so.FindProperty(attr.minVariable).intValue;
+                        int maxValue = so.FindProperty(attr.maxVariable).intValue;
+                        EditorGUILayout.IntSlider(prop, minValue, maxValue, new GUIContent(prop.displayName));
+                    }
                 }
                 else
                 {

diff --git a/Runtime/LLM.cs b/Runtime/LLM.cs
@@ -37,7 +37,7 @@ public class LLM : MonoBehaviour
         [LLMAdvanced] public bool dontDestroyOnLoad = true;
         /// <summary> Size of the prompt context (0 = context size of the model).
         /// This is the number of tokens the model can take as input when generating responses. </summary>
-        [ModelAdvanced] public int contextSize = 0;
+        [DynamicRange("minContextLength", "maxContextLength", false), Model] public int contextSize = 8192;
         /// <summary> Batch size for prompt processing. </summary>
         [ModelAdvanced] public int batchSize = 512;
         /// <summary> a base prompt to use as a base for all LLMCharacter objects </summary>
@@ -74,9 +74,10 @@ public class LLM : MonoBehaviour
         [SerializeField]
         private string SSLKey = "";
         public string SSLKeyPath = "";
-        public static int maxContextLength = 32768;
 
         /// \cond HIDE
+        public int minContextLength = 0;
+        public int maxContextLength = 0;
 
         IntPtr LLMObject = IntPtr.Zero;
         List<LLMCharacter> clients = new List<LLMCharacter>();
@@ -212,16 +213,9 @@ public void SetModel(string path)
                 ModelEntry modelEntry = LLMManager.Get(model);
                 if (modelEntry == null) modelEntry = new ModelEntry(GetLLMManagerAssetRuntime(model));
                 SetTemplate(modelEntry.chatTemplate);
-                if (contextSize > modelEntry.contextLength)
-                {
-                    contextSize = 0;
-                    LLMUnitySetup.LogWarning($"The selected context size {contextSize} is larger than the context size supported from the {path} model ({modelEntry.contextLength}), resetting it to use the model's default");
-                }
-                if (contextSize == 0 && modelEntry.contextLength > maxContextLength)
-                {
-                    contextSize = maxContextLength;
-                    LLMUnitySetup.LogWarning($"The {path} model has very large context size ({modelEntry.contextLength}), it was automatically set to {maxContextLength} to avoid filling up the RAM");
-                }
+
+                maxContextLength = modelEntry.contextLength;
+                if (contextSize > maxContextLength) contextSize = maxContextLength;
             }
 #if UNITY_EDITOR
             if (!EditorApplication.isPlaying) EditorUtility.SetDirty(this);

diff --git a/Runtime/LLMUnitySetup.cs b/Runtime/LLMUnitySetup.cs
@@ -38,6 +38,21 @@ public IntAttribute(int min, int max)
         }
     }
 
+    [AttributeUsage(AttributeTargets.Field, Inherited = true, AllowMultiple = false)]
+    public class DynamicRangeAttribute : PropertyAttribute
+    {
+        public readonly string minVariable;
+        public readonly string maxVariable;
+        public bool intOrFloat;
+
+        public DynamicRangeAttribute(string minVariable, string maxVariable, bool intOrFloat)
+        {
+            this.minVariable = minVariable;
+            this.maxVariable = maxVariable;
+            this.intOrFloat = intOrFloat;
+        }
+    }
+
     public class LLMAttribute : PropertyAttribute {}
     public class LLMAdvancedAttribute : PropertyAttribute {}
     public class LocalRemoteAttribute : PropertyAttribute {}