Views
No views yet
1from transformers import AutoTokenizer, EncoderDecoderModel
2
3# Load the CodeRosetta model and tokenizer
4model = EncoderDecoderModel.from_pretrained('CodeRosetta/CodeRosetta_cpp2cuda_ft')
5tokenizer = AutoTokenizer.from_pretrained('CodeRosetta/CodeRosetta_cpp2cuda_ft')
6
7# Encode the input C++ Code
8input_cpp_code = "void add_100 ( int numElements , int * data ) { for ( int idx = 0 ; idx < numElements ; idx ++ ) { data [ idx ] += 100 ; } }"
9input_ids = tokenizer.encode(input_cpp_code, return_tensors="pt")
10
11# Set the start token to <CUDA>
12start_token = "<CUDA>"
13decoder_start_token_id = tokenizer.convert_tokens_to_ids(start_token)
14
15# Generate the CUDA code
16output = model.generate(
17 input_ids=input_ids,
18 decoder_start_token_id=decoder_start_token_id,
19 max_length=256
20)
21
22# Decode and print the generated output
23generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
24print(generated_code)1@inproceedings{coderosetta:neurips:2024,
2 title = {CodeRosetta: Pushing the Boundaries of Unsupervised Code Translation for Parallel Programming},
3 author = {TehraniJamsaz, Ali and Bhattacharjee, Arijit and Chen, Le and Ahmed, Nesreen K and Yazdanbakhsh, Amir and Jannesari, Ali},
4 booktitle = {NeurIPS},
5 year = {2024},
6}