Update endpoints.py

endpoints.py

@@ -16,7 +16,10 @@ import requests
 def LLM(llm_name, length):
     print(llm_name)
     tokenizer = AutoTokenizer.from_pretrained(llm_name)
-    model = AutoModelForCausalLM.from_pretrained(llm_name)
+    model = AutoModelForCausalLM.from_pretrained(llm_name,
+                                                 trust_remote_code=True, 
+                                                 device_map="auto",
+                                                 load_in_8bit=True)
     pipe = pipeline("text-generation",
                     model=model,
                     tokenizer=tokenizer,
@@ -28,8 +31,7 @@ def LLM(llm_name, length):
     return pipe
 
 
-pipe = LLM("codellama/CodeLlama-7b-hf",4000)
-
+pipe = LLM("replit/replit-code-v1-3b",4000)
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # Mistral 7B