Update README.md

README.md

@@ -11,27 +11,7 @@ To use notdiamond-0001, format your queries using the following prompt with your
 ``` python
     query = "Can you write a function that counts from 1 to 10?"
 
-    formatted_prompt = f"""I would like you to determine whether a given query should be sent to GPT-3.5 or GPT-4.
-        In general, the following types of queries should get sent to GPT-3.5:
-        Explanation
-        Summarization
-        Writing
-        Informal conversation
-        History, government, economics, literature, social studies
-        Simple math questions
-        Simple coding questions
-        Simple science questions
-
-        In general, the following types of queries should get sent to GPT-4:
-        Advanced coding questions
-        Advanced math questions
-        Advanced science questions
-        Legal questions
-        Medical questions
-        Sensitive/inappropriate queries
-
-        Your job is to determine whether the following query should be sent to GPT-3.5 or GPT-4.
-
+    formatted_prompt = f"""Determine whether the following query should be sent to GPT-3.5 or GPT-4.
         Query:
         {query}"""
 ```
@@ -45,9 +25,8 @@ You can then determine the model to call as follows
     id2label = {0: 'gpt-3.5', 1: 'gpt-4'}
     tokenizer = AutoTokenizer.from_pretrained("notdiamond/notdiamond-0001")
     model = AutoModelForSequenceClassification.from_pretrained("notdiamond/notdiamond-0001")
-    max_length = self._get_max_length(model)
 
-    inputs = tokenizer(formatted_prompt, truncation=True, max_length=max_length, return_tensors="pt")
+    inputs = tokenizer(formatted_prompt, truncation=True, max_length=512, return_tensors="pt")
     logits = model(**inputs).logits
     model_id = logits.argmax().item()
     model_to_call = id2label[model_id]