Update README.md

README.md

@@ -1,5 +1,11 @@
 ---
 license: apache-2.0
+datasets:
+- bigcode/the-stack
+- bigcode/the-stack-v2
+- bigcode/starcoderdata
+- bigcode/commitpack
+library_name: transformers
 ---
 
 # Model Description
@@ -109,16 +115,13 @@ write_dist_matrix(dist_matrix, max_dist, output_data_filename_, verbose=True)
 end = time.perf_counter()
 """
 
-tokenizer = AutoTokenizer.from_pretrained('mellum-base-4b')
-model = AutoModelForCausalLM.from_pretrained('mellum-base-4b')
+tokenizer = AutoTokenizer.from_pretrained('JetBrains/Mellum-4b-base')
+model = AutoModelForCausalLM.from_pretrained('JetBrains/Mellum-4b-base')
 encoded_input = tokenizer(example, return_tensors='pt', return_token_type_ids=False)
 input_len = len(encoded_input["input_ids"][0])
 out = model.generate(
     **encoded_input,
     max_new_tokens=100,
-    num_beams=1,
-    pad_token_id=tokenizer.eos_token_id,
-    eos_token_id=tokenizer.eos_token_id,
 )
 print("### Context")
 print(tokenizer.decode(out[0][:input_len]))
@@ -138,4 +141,4 @@ If you use this model, please cite:
 ```
 
 # Contact
-For questions, collaborations and requests reach us out via [email protected]
+For questions, collaborations and requests reach us out via [email protected]