|
|
--- |
|
|
license: mit |
|
|
--- |
|
|
This is a stable version of BSJCode, a model capable of fixing as well as optimizing java code. |
|
|
|
|
|
------------------ |
|
|
```bash |
|
|
## How to use it: |
|
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig |
|
|
import torch |
|
|
|
|
|
# Load the model and tokenizer |
|
|
bnb_config = BitsAndBytesConfig( |
|
|
load_in_4bit=True, |
|
|
bnb_4bit_quant_type="nf4", |
|
|
bnb_4bit_compute_dtype=torch.float16, |
|
|
bnb_4bit_use_double_quant=True, |
|
|
llm_int8_enable_fp32_cpu_offload=True |
|
|
) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
"BSAtlas/BSJCode-1-Stable", |
|
|
quantization_config=bnb_config, |
|
|
device_map="auto" |
|
|
).to(device="cuda") |
|
|
tokenizer = AutoTokenizer.from_pretrained("BSAtlas/BSJCode-1-Stable") |
|
|
|
|
|
|
|
|
def detect_and_fix_bugs(code_snippet): |
|
|
# Prepare the prompt |
|
|
prompt = f"""You are an expert Java code optimizer and bug fixer. |
|
|
Analyze the following code, identify any bugs or inefficiencies, |
|
|
and provide an optimized and corrected version: |
|
|
|
|
|
```java |
|
|
{code_snippet} |
|
|
``` |
|
|
|
|
|
Optimized and Fixed Code:""" |
|
|
|
|
|
# Tokenize the input |
|
|
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024) |
|
|
|
|
|
# Generate code |
|
|
with torch.no_grad(): |
|
|
outputs = model.generate( |
|
|
input_ids=inputs['input_ids'], |
|
|
attention_mask=inputs['attention_mask'], |
|
|
max_length=1024, # Adjust based on your model's training |
|
|
num_return_sequences=1, |
|
|
do_sample=True, |
|
|
temperature=0.7, |
|
|
top_k=50, |
|
|
top_p=0.95, |
|
|
repetition_penalty=1.2 |
|
|
) |
|
|
|
|
|
# Decode the output |
|
|
generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
|
|
# Extract the code portion after the prompt |
|
|
code_start = generated_code.find("Optimized and Fixed Code:") |
|
|
if code_start != -1: |
|
|
fixed_code = generated_code[code_start + len("Optimized and Fixed Code:"):].strip() |
|
|
else: |
|
|
fixed_code = generated_code |
|
|
|
|
|
return fixed_code |
|
|
|
|
|
|
|
|
sample_code = """ |
|
|
public class ThreadSafetyExample { |
|
|
private int counter = 0; |
|
|
|
|
|
public void increment() { |
|
|
// Not thread-safe method |
|
|
counter++; |
|
|
} |
|
|
|
|
|
public int getCounter() { |
|
|
return counter; |
|
|
} |
|
|
|
|
|
public static void main(String[] args) { |
|
|
ThreadSafetyExample example = new ThreadSafetyExample(); |
|
|
|
|
|
Thread t1 = new Thread(() -> { |
|
|
for (int i = 0; i < 1000; i++) { |
|
|
example.increment(); |
|
|
} |
|
|
}); |
|
|
|
|
|
Thread t2 = new Thread(() -> { |
|
|
for (int i = 0; i < 1000; i++) { |
|
|
example.increment(); |
|
|
} |
|
|
}); |
|
|
|
|
|
t1.start(); |
|
|
t2.start(); |
|
|
|
|
|
try { |
|
|
t1.join(); |
|
|
t2.join(); |
|
|
} catch (InterruptedException e) { |
|
|
e.printStackTrace(); |
|
|
} |
|
|
|
|
|
System.out.println("Final Counter: " + example.getCounter()); |
|
|
} |
|
|
} |
|
|
""" |
|
|
fixed_code = detect_and_fix_bugs(sample_code) |
|
|
print(fixed_code) |
|
|
-------------------------------------------------------------------------- |
|
|
|
|
|
|
