Update endpoints.py

endpoints.py

@@ -13,25 +13,33 @@ 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
-
-
-pipe = LLM("replit/replit-code-v1-3b",4000)
+# 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
 # tokenizer = AutoTokenizer.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # base_model = AutoModelForCausalLM.from_pretrained("WizardLM/WizardCoder-1B-V1.0")
 # Mistral 7B