From 16c755257b345e78bfefac2a513f744339d04013 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 30 Sep 2023 15:37:30 -0700
Subject: [PATCH] add support for custom hf prompt templates

---
 litellm/__pycache__/main.cpython-311.pyc  | Bin 52139 -> 52158 bytes
 litellm/__pycache__/utils.cpython-311.pyc | Bin 136629 -> 137059 bytes
 litellm/llms/prompt_templates/factory.py  | 122 ++++++++++++++++++----
 litellm/tests/test_hf_prompt_templates.py |  27 +++--
 litellm/tests/test_streaming.py           | 106 ++++++++++---------
 pyproject.toml                            |   2 +-
 6 files changed, 172 insertions(+), 85 deletions(-)

diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc
index 61a2e4b71bfad0a5aa50d74b8dee28ce82d71601..9f0a8d2cecaa0747921fc2b41286a5c0e7baf816 100644
GIT binary patch
delta 880
zcmXAkZAep57=X`v?{w2`ZgabLGdFFb`GMxgEH#^Qx82Y`m1coN$xJISBts)h5y4D@
z(zt{uLK;Mc8N?e|(AR!Qh=HL&VLu9sLL$6VNs1ZWdptkR^SmGDyyr{|u`h>M<6WaM
zmMAOljH9t>+{nW^#_@|<IEBlu1qL?p#r8FpwKOw`sUkof{z=OPc!Qbgi2x-iX7~Yy
zus_qyvc~2poXpMvXvR-D%K+^-o!buZTq@5;sK>Q@<@_rw@n```Q7Ck4APY|vKVqQ@
zxxF5MDXG#s3sn+VhOi;cluf`mwv{h4kd1Y|CV=zOtUp~_-!RhYXq}U~Jtfkhw*f{Z
zccBaK1!}fID{dYSYvCHMoD2Y5!ot&yPQA<%Ge6jRwWF^?YPgu9sd!giD~)zLq_Y>a
zxYRQSNgSEZspNzpN9I_+UFCdY#vu(^U_l9_v0leMj=awmQ`OR<O-xkzInPI1mvla(
z8h_(`Zk6AQD>78BxCLINT<mP&NZ3ZN`q%5|;>fDUpRAUCTZIIb8#%rdm4D4XR%OIS
zZ;I^Jbd-|F9`z)Nlv+9s?4}_=qmU5iFILkbQL7O|gF^jQK{Qwd(Zz;Dmp_I!RHTb}
zRQ5ZBG)0b((1Eh8xuGhDK))TT-sv;Zrkb*R%ENgg{YB_#g0PEbj+p{cRDxWr5G-Ot
z!DOotEKpdW<csj<m0^Hec)5Q;17EQAx*gybwhY_>xR3T58XcG<=kQxt#pU5X9dLN=
zHlqhWwoQ&j!&BrR-(+AOA3Vtd_=+E<9s>Np-e&{R5Jvv>sIE&iL##~H{i8lkW@#8r
zOjcP-RY%M=IfiO1vDsxkRRf8)%SNhk#9)z4RLw-Emn~G|iP0iksq)0a%XX>`Vt2|3
zRGq}a$u80CRi-BqWpXl2?;z1uIh`tf)7HzvT83fPiL!-R!svQ<Sw~wAV)Klb`45fW
B2=)K~

delta 929
zcmXAlTWAwO6hP<hq}epDO`1*8nnq2jrj^Dtrb>M@x=9w5LZw;-QIX&SK@r6VDn6lV
z{qeyfS{1E=D2P(MrL`Yw_rniS5KTTx{SZXOis-~AYJBWY_hZhvckZ1#v-i)?8)vBN
zn#*M;#@DqcxNiPsR{-9jJ1`<b4oAaz0e;}UiC-x7Y%U?zssPLJ=cG7*j`h(HKoZru
z48SQoRv)3%x%nU7Y={ADz!$MmfNl6Nz8j!lo0~vbiC+>%{~*pxiU8A5X^L2&7S}Z2
zrqGK1IZ1%q+QOE{(5m@U2u1Bt>IR(09dkzon2f8_8v(Xxk26tm<=S)mDz`t?4kc%5
zCtG}4dnQV)dqwS1EClDYNK+X5vK`}$;pMA~A{@cr>$3oRaNdS>)zNPO;t3X{nx51%
zX<9*=#z>iKV(-o&`Of~{{cl23hvk|c_3ke1;=vGp?7RjCB=Rt}(68hq@|0$RCZC7u
zmRLxhf*Is3_70Ke-Ax?=EeK_X7sRN!QHZOxCP!p7Wb$QWns>jlrCnzHI*^W-eA~RB
z&g7?v5-`b{tGD_|QRb%?{hEB#Z`|Q?GBd&4{q9q$O#X<aYfS#~q^Zf0%nvKgu<V^p
zO0#%#X|{Ktlm2mLd88}uOtV?4tD;3wodz9a8>-W*s9~zB;Y>L%SP)ekxwKU(lMGo#
z^hK$$xq5r0!heRocxqbWMIEUG=h8-%KSW<qrNnt=Re83IOHQ&Yxdhd7RX!!xV6ef+
zC-CplZh(t;_E_EmFLBfHAixLgJaHA^I!-xdv4LBQcRzt5mb$xaAmiaHf&<dnb+g9`
z16bd8T7YNxpnnR$D||n23*aqw51z0B!}_5MwgXuY*cma`VTYSh9$JayV<J}@@yM*4
ztDVR}=HTii6+!0W>LyMvlel__&B44}D~QX>d|U&>8(=}Ml_Xfrs<>7Yub+jpo)%;J
qcw$VRz|)h7)yJY-`JKjYRulz6_)3hWNQq$U3@{rnX~RR0T7>_!$`;}P

diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc
index 338c3b089645d248ad6f355d849ed2895e4a4341..d3a54a2d4ce6f9c1df91226ee7d4874596cac59a 100644
GIT binary patch
delta 837
zcmZutYe<t}6#ky~+qXGsGnuO`FY~%UyWh-^&_XH;T#KlbiqQ%ZSr+~1ht`4x`okzM
zogfOSAaQ2i=BrlIif)o^lNm)9jV=aOSlJDV`o8>A(D`wm%XwZ7=k<9bUNlGOFX(iU
zf?V@TOglZE!+M_{{f%s5YR!CW0)1=5J9N-%FNQ?5Yc=_s5Wol>y@&x(avh*;%_x`2
zQhL@b4}#Qu2}e++#A2zS1#P0rm8OJZ$<hiZUXiC2qoU|4C084c%fw1=+Yv_z?bwYM
zRMCzdXjQC$$dq**xQ0=3`(P9g?a!oaFH-4#C&H>@aBbz)orp$MxJ)Q!$^}X}<9U52
z96$n<bYU@Nc45&Zuahig%$Z2)ZdkO*yM;1#9!EiaUYKcBH<scCxx3+;3@@VM9?V1%
zRreqRDdg`#o2m|i!1r9n31r!{g5VaOnyn4auGA~(L33WnoEJ9dRp`grMy5T0akh@>
z(npdr8tnC!WBP}N>0v|0xKK{Zui_c@QoJ8~P)t>RxZtE8eyoIpa(ZP|q(Bk{dy#_8
zlL0&5(I=n5a==es!y810Vp1O(&0*v07So_HD`?E(Kd&Q-X}749t=~QulJO?^q3v<P
zz{1>^0<93xFj)k&kv5YUFf6kf#Gp<sJ3)gaJ3*6C;d89(rwO4LB^jC)Y1fIN#AJJf
zILMUbV1y()gSz|-Sv<*FF-3-9wL$4Klo!KsGaQ-X@RHmmjs@cLS&)VFe1YW1B-<_Y
zql(Ag#xEE*F-iU$;UCz?=ibE@pq5wmV*)tE58lHXE$-5tA&lWA?|6XC42@j(2=_Eh
zqvMZ95T!vP9~*^RjeK769*HWH^Qw>71XS`@pKu1K;hxWk2hP#iFF1_`p7RwO7@BzP
ocdXT5z&FB*RG7lECH7uFuT(<XxbXk};g3J@_jDI8jA2Fi4Vg<CUH||9

delta 752
zcmZutT}YE*6#ky~+uSB@m2)aH-CU&NUw>v~5osWak!F;ds8D**&BB}LLQA1YA*@h2
zL?&KYXsw*to355C`V(zBj6|ABfn`6kKxtr9@0T|Py%*>CInUv|oc9iGe~;F1U9S(7
z#0z#sRWCI(82Swu?q;*AS~exbQLr0dkxX+vcrB~#ZgSs*7l*0&7Di>2eSs<*s1w9N
z(+<(dCC6=?gRI12sNfEqvTVPqgks3jiz<}S*<O5*nSGAzPFxnGoxb#8EyeZWI7+Fp
z566(GSY8>*y1Ve8oa*~wA{WOGOz_rwSOensQGr;hAHY`1A3(+mZ&SzBvMLLi2Vp@x
zWeuVf+o^RB{VV+pI_*LttkmT~3T)(a!O5E*pbgBWC950e*gU2?g#g>l^vScflp95z
zmP>|7W7LE(B_P#Mu^ZE<rp!k;feO-lU`Gj=hOh^RsB}oE0tMn}ehBNZZzYh+YdzwF
z$l%V$_z0f<6c?b&Rw*ljOF8u98Tgs<aA|Z)c0_?r@@kk6UR`KmwCpuz6&mF+eTZ<z
zjB4SGMVl0!o^T*i8jnzt<B_37dU@PxE7Hn7rX+n@wQzj;z41a!ZO91|@OGzBRR3ya
z!XIv~Oq2bo*_o9|@{A2asArPYm1*?zIfl`|U%tREnDUugTtfN>j&T35xG5U>wh=4>
z7kTO{<mvE`PQJlsO!A|XC}HT}t|^RY*m4+edWTgSMDQ!~s1Ly|ZVtezLJi-sh+?3Q
zw=AIpIL|kKM=U^87{oQS^MxP|GIa8sW#nt%;mwSdsUY!BA#B00u}Y2g0qOts+aG`N
M@9uuCjbvr`4c#RFeE<Le

diff --git a/litellm/llms/prompt_templates/factory.py b/litellm/llms/prompt_templates/factory.py
index d47a7486db3..381e8b8fdc3 100644
--- a/litellm/llms/prompt_templates/factory.py
+++ b/litellm/llms/prompt_templates/factory.py
@@ -1,3 +1,7 @@
+import requests, traceback
+import json
+from jinja2 import Template, exceptions, Environment, meta
+
 def default_pt(messages):
     return " ".join(message["content"] for message in messages)
 
@@ -104,6 +108,76 @@ def phind_codellama_pt(messages):
             prompt += "### Assistant\n" + message["content"] + "\n\n"
     return prompt
 
+def hf_chat_template(model: str, messages: list):
+    ## get the tokenizer config from huggingface
+    def _get_tokenizer_config(hf_model_name):
+        url = f"https://huggingface.co/{hf_model_name}/raw/main/tokenizer_config.json"
+        # Make a GET request to fetch the JSON data
+        response = requests.get(url)
+        if response.status_code == 200:
+            # Parse the JSON data
+            tokenizer_config = json.loads(response.content)
+            return {"status": "success", "tokenizer": tokenizer_config}
+        else:
+            return {"status": "failure"}
+    tokenizer_config = _get_tokenizer_config(model)
+    if tokenizer_config["status"] == "failure" or "chat_template" not in tokenizer_config["tokenizer"]:
+        raise Exception("No chat template found")
+    ## read the bos token, eos token and chat template from the json 
+    tokenizer_config = tokenizer_config["tokenizer"]
+    bos_token = tokenizer_config["bos_token"]
+    eos_token = tokenizer_config["eos_token"]
+    chat_template = tokenizer_config["chat_template"]
+
+    def raise_exception(message):
+        raise Exception(f"Error message - {message}")
+    
+    # Create a template object from the template text
+    env = Environment()
+    env.globals['raise_exception'] = raise_exception
+    template = env.from_string(chat_template)
+
+    def _is_system_in_template():
+        try:
+            # Try rendering the template with a system message
+            response = template.render(messages=[{"role": "system", "content": "test"}], eos_token= "<eos>", bos_token= "<bos>")
+            return True
+
+        # This will be raised if Jinja attempts to render the system message and it can't
+        except:
+            return False
+        
+    try: 
+        # Render the template with the provided values
+        if _is_system_in_template(): 
+            rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=messages)
+        else: 
+            # treat a system message as a user message, if system not in template
+            try:
+                reformatted_messages = []
+                for message in messages: 
+                    if message["role"] == "system": 
+                        reformatted_messages.append({"role": "user", "content": message["content"]})
+                    else:
+                        reformatted_messages.append(message)
+                rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=reformatted_messages)
+            except Exception as e:
+                if "Conversation roles must alternate user/assistant" in str(e): 
+                    # reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, add a blank 'user' or 'assistant' message to ensure compatibility
+                    new_messages = []
+                    for i in range(len(reformatted_messages)-1): 
+                        new_messages.append(reformatted_messages[i])
+                        if reformatted_messages[i]["role"] == reformatted_messages[i+1]["role"]:
+                            if reformatted_messages[i]["role"] == "user":
+                                new_messages.append({"role": "assistant", "content": ""})
+                            else:
+                                new_messages.append({"role": "user", "content": ""})
+                    new_messages.append(reformatted_messages[-1])
+                    rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=new_messages)
+        return rendered_text
+    except: 
+        raise Exception("Error rendering template")
+
 # Custom prompt template
 def custom_prompt(role_dict: dict, messages: list, initial_prompt_value: str="", final_prompt_value: str=""):
     prompt = initial_prompt_value
@@ -117,27 +191,31 @@ def custom_prompt(role_dict: dict, messages: list, initial_prompt_value: str="",
     return prompt
 
 def prompt_factory(model: str, messages: list):
+    original_model_name = model
     model = model.lower()
-    if "meta-llama/llama-2" in model:
-        if "chat" in model:
+    try:
+        if "meta-llama/llama-2" in model:
+            if "chat" in model:
+                return llama_2_chat_pt(messages=messages)
+        elif "tiiuae/falcon" in model: # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.
+            if model == "tiiuae/falcon-180B-chat":
+                return falcon_chat_pt(messages=messages)
+            elif "instruct" in model:
+                return falcon_instruct_pt(messages=messages)
+        elif "mosaicml/mpt" in model:
+            if "chat" in model:
+                return mpt_chat_pt(messages=messages)
+        elif "codellama/codellama" in model:
+            if "instruct" in model:
+                return llama_2_chat_pt(messages=messages) # https://huggingface.co/blog/codellama#conversational-instructions
+        elif "wizardlm/wizardcoder" in model:
+            return wizardcoder_pt(messages=messages)
+        elif "phind/phind-codellama" in model:
+            return phind_codellama_pt(messages=messages)
+        elif "togethercomputer/llama-2" in model and ("instruct" in model or "chat" in model):
             return llama_2_chat_pt(messages=messages)
-    elif "tiiuae/falcon" in model: # Note: for the instruct models, it's best to use a User: .., Assistant:.. approach in your prompt template.
-        if model == "tiiuae/falcon-180B-chat":
-            return falcon_chat_pt(messages=messages)
-        elif "instruct" in model:
-            return falcon_instruct_pt(messages=messages)
-    elif "mosaicml/mpt" in model:
-        if "chat" in model:
-            return mpt_chat_pt(messages=messages)
-    elif "codellama/codellama" in model:
-        if "instruct" in model:
-            return llama_2_chat_pt(messages=messages) # https://huggingface.co/blog/codellama#conversational-instructions
-    elif "wizardlm/wizardcoder" in model:
-        return wizardcoder_pt(messages=messages)
-    elif "phind/phind-codellama" in model:
-        return phind_codellama_pt(messages=messages)
-    elif "togethercomputer/llama-2" in model and ("instruct" in model or "chat" in model):
-        return llama_2_chat_pt(messages=messages)
-    elif "mistralai/mistral" in model and "instruct" in model: 
-        return mistral_instruct_pt(messages=messages)
-    return default_pt(messages=messages) # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2)
\ No newline at end of file
+        else: 
+            return hf_chat_template(original_model_name, messages)
+    except:
+        traceback.print_exc()
+        return default_pt(messages=messages) # default that covers Bloom, T-5, any non-chat tuned model (e.g. base Llama2)
\ No newline at end of file
diff --git a/litellm/tests/test_hf_prompt_templates.py b/litellm/tests/test_hf_prompt_templates.py
index 0066a9e9d7c..9f1aedcff3a 100644
--- a/litellm/tests/test_hf_prompt_templates.py
+++ b/litellm/tests/test_hf_prompt_templates.py
@@ -1,17 +1,22 @@
-# import sys, os
-# import traceback
-# from dotenv import load_dotenv
+import sys, os
+import traceback
+from dotenv import load_dotenv
 
-# load_dotenv()
-# import os
+load_dotenv()
+import os
 
-# sys.path.insert(
-#     0, os.path.abspath("../..")
-# )  # Adds the parent directory to the system path
-# import pytest
-# import litellm
-# from litellm import embedding, completion, text_completion
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest
+from litellm.llms.prompt_templates.factory import prompt_factory
 
+def test_prompt_formatting(): 
+    try:
+        prompt = prompt_factory(model="mistralai/Mistral-7B-Instruct-v0.1", messages=[{"role": "system", "content": "Be a good bot"}, {"role": "user", "content": "Hello world"}])
+        assert prompt == "<s>[INST] Be a good bot [/INST]</s> [INST] Hello world [/INST]"
+    except Exception as e: 
+        pytest.fail(f"An exception occurred: {str(e)}")
 # def logger_fn(user_model_dict):
 #     return
 #     print(f"user_model_dict: {user_model_dict}")
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index 78a871466e6..638253cda1b 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -314,58 +314,62 @@ def test_completion_cohere_stream_bad_key():
 
 # test_completion_nlp_cloud_bad_key()
 
-# def test_completion_hf_stream():
-#     try:
-#         litellm.set_verbose = True
-#         # messages = [
-#         #     {
-#         #         "content": "Hello! How are you today?",
-#         #         "role": "user"
-#         #     },
-#         # ]
-#         # response = completion(
-#         #     model="huggingface/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud", stream=True, max_tokens=1000
-#         # )
-#         # complete_response = ""
-#         # # Add any assertions here to check the response
-#         # for idx, chunk in enumerate(response):
-#         #     chunk, finished = streaming_format_tests(idx, chunk)
-#         #     if finished:
-#         #         break
-#         #     complete_response += chunk
-#         # if complete_response.strip() == "": 
-#         #     raise Exception("Empty response received")
-#         # completion_response_1 = complete_response
-#         messages = [
-#             {
-#                 "content": "Hello! How are you today?",
-#                 "role": "user"
-#             },
-#             {
-#                 "content": "I'm doing well, thank you for asking! I'm excited to be here and help you with any questions or concerns you may have. What can I assist you with today?</s>",
-#                 "role": "assistant"
-#             },
-#         ]
-#         response = completion(
-#             model="huggingface/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud", stream=True, max_tokens=1000
-#         )
-#         complete_response = ""
-#         # Add any assertions here to check the response
-#         for idx, chunk in enumerate(response):
-#             chunk, finished = streaming_format_tests(idx, chunk)
-#             if finished:
-#                 break
-#             complete_response += chunk
-#         if complete_response.strip() == "": 
-#             raise Exception("Empty response received")
-#         # print(f"completion_response_1: {completion_response_1}")
-#         print(f"completion_response: {complete_response}")
-#     except InvalidRequestError as e:
-#         pass
-#     except Exception as e:
-#         pytest.fail(f"Error occurred: {e}")
+def test_completion_hf_stream():
+    try:
+        litellm.set_verbose = True
+        # messages = [
+        #     {
+        #         "content": "Hello! How are you today?",
+        #         "role": "user"
+        #     },
+        # ]
+        # response = completion(
+        #     model="huggingface/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud", stream=True, max_tokens=1000
+        # )
+        # complete_response = ""
+        # # Add any assertions here to check the response
+        # for idx, chunk in enumerate(response):
+        #     chunk, finished = streaming_format_tests(idx, chunk)
+        #     if finished:
+        #         break
+        #     complete_response += chunk
+        # if complete_response.strip() == "": 
+        #     raise Exception("Empty response received")
+        # completion_response_1 = complete_response
+        messages = [
+            {
+                "content": "Hello! How are you today?",
+                "role": "user"
+            },
+            {
+                "content": "I'm doing well, thank you for asking! I'm excited to be here and help you with any questions or concerns you may have. What can I assist you with today?",
+                "role": "assistant"
+            },
+            {
+                "content": "What is the price of crude oil?",
+                "role": "user"
+            },
+        ]
+        response = completion(
+            model="huggingface/mistralai/Mistral-7B-Instruct-v0.1", messages=messages, api_base="https://n9ox93a8sv5ihsow.us-east-1.aws.endpoints.huggingface.cloud", stream=True, max_tokens=1000
+        )
+        complete_response = ""
+        # Add any assertions here to check the response
+        for idx, chunk in enumerate(response):
+            chunk, finished = streaming_format_tests(idx, chunk)
+            if finished:
+                break
+            complete_response += chunk
+        if complete_response.strip() == "": 
+            raise Exception("Empty response received")
+        # print(f"completion_response_1: {completion_response_1}")
+        print(f"completion_response: {complete_response}")
+    except InvalidRequestError as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
 
-# test_completion_hf_stream()
+test_completion_hf_stream()
 
 # def test_completion_hf_stream_bad_key():
 #     try:
diff --git a/pyproject.toml b/pyproject.toml
index 13a808dd08d..545301001ed 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.802"
+version = "0.1.803"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"