Update endpoints.py

endpoints.py

@@ -13,24 +13,22 @@ import requests
 # 	response = requests.post(API_URL, headers=headers, json=payload)
 # 	return response.json()
 
-# def LLM(llm_name, length):
-#     print(llm_name)
-#     tokenizer = AutoTokenizer.from_pretrained(llm_name)
-#     model = AutoModelForCausalLM.from_pretrained(llm_name)
-#     pipe = pipeline("text-generation",
-#                     model=model,
-#                     tokenizer=tokenizer,
-#                     max_length=length,
-#                     do_sample=True,
-#                     top_p=0.95,
-#                     repetition_penalty=1.2,
-#                    )
-#     return pipe
-# Load model directly
-# Use a pipeline as a high-level helper
-from transformers import pipeline
-
-pipe = pipeline("text-generation", model="codellama/CodeLlama-7b-hf")
+def LLM(llm_name, length):
+    print(llm_name)
+    tokenizer = AutoTokenizer.from_pretrained(llm_name)
+    model = AutoModelForCausalLM.from_pretrained(llm_name)
+    pipe = pipeline("text-generation",
+                    model=model,
+                    tokenizer=tokenizer,
+                    max_length=length,
+                    do_sample=True,
+                    top_p=0.95,
+                    repetition_penalty=1.2,
+                   )
+    return pipe
+
+
+pipe = LLM("codellama/CodeLlama-7b-hf",4000)
 
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")