Update endpoints.py

endpoints.py

@@ -28,23 +28,24 @@ def LLM(llm_name, length):
     return pipe
 
 
-pipe = LLM("microsoft/phi-2",2000)
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # Mistral 7B
 # mistral_llm = LLM("mistralai/Mistral-7B-v0.1",30000)
-mistral_llm = pipe
+mistral_llm = LLM("microsoft/phi-2",2000)
 
 # WizardCoder 13B
 # wizard_llm = LLM("WizardLM/WizardCoder-Python-13B-V1.0",8000)
-wizard_llm = pipe
+wizard_llm =  LLM("WizardLM/WizardCoder-3B-V1.0",4000)
 # hf_llm = HuggingFacePipeline(pipeline=pipe)
 
 def ask_model(model, prompt):
     if(model == 'mistral'):
-        return mistral_llm(prompt)
+        result = mistral_llm(prompt)
+        return result
     if(model == 'wizard'):
-        return wizard_llm(prompt)
+         result =  wizard_llm(prompt)
+         return result