Update endpoints.py

endpoints.py

@@ -13,33 +13,25 @@ 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,
-#                                                  trust_remote_code=True, 
-#                                                  device_map="auto",
-#                                                  load_in_8bit=True)
-#     pipe = pipeline("text-generation",
-#                     model=model,
-#                     tokenizer=tokenizer,
-#                     max_length=length,
-#                     do_sample=True,
-#                     top_p=0.95,
-#                     repetition_penalty=1.2,
-#                    )
-#     return pipe
-tokenizer = AutoTokenizer.from_pretrained('replit/replit-code-v1_5-3b', trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained('replit/replit-code-v1_5-3b', trust_remote_code=True)
-
-x = tokenizer.encode('def fibonacci(n): ', return_tensors='pt')
-y = model.generate(x, max_length=100, do_sample=True, top_p=0.95, top_k=4, temperature=0.2, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
-
-# decoding
-generated_code = tokenizer.decode(y[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)
-print(generated_code)
-
-pipe = generated_code
+def LLM(llm_name, length):
+    print(llm_name)
+    tokenizer = AutoTokenizer.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,
+                    max_length=length,
+                    do_sample=True,
+                    top_p=0.95,
+                    repetition_penalty=1.2,
+                   )
+    return pipe
+
+
+pipe = LLM("WizardLM/WizardCoder-Python-7B-V1.0",4000)
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # Mistral 7B