Update README.md

README.md

@@ -10,24 +10,31 @@ This is a repository for fine-tuned CodeLlama-7b model in the Hugging Face Trans
 # Model use
 
 ```python
-  from transformers import AutoModelForCausalLM, AutoTokenizer
-
-  # Load pre-trained model and tokenizer
-  model_name = 'JetBrains/CodeLlama-7B-Kexer'  # Replace with the desired model name
-  tokenizer = AutoTokenizer.from_pretrained(model_name)
-  model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
-
-  # Encode input text
-  input_text = """This function takes an integer n and returns factorial of a number:
-  fun factorial(n: Int): Int {"""
-  input_ids = tokenizer.encode(input_text, return_tensors='pt').to('cuda')
-
-  # Generate text
-  output = model.generate(input_ids, max_length=150, num_return_sequences=1, no_repeat_ngram_size=2, early_stopping=True)
-
-  # Decode and print the generated text
-  generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-  print(generated_text)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load pre-trained model and tokenizer
+model_name = 'JetBrains/CodeLlama-7B-Kexer'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to('cuda')
+
+# Create and encode input
+input_text = """\
+This function takes an integer n and returns factorial of a number:
+fun factorial(n: Int): Int {\
+"""
+input_ids = tokenizer.encode(
+    input_text, return_tensors='pt'
+).to('cuda')
+
+# Generate
+output = model.generate(
+    input_ids, max_length=150, num_return_sequences=1, 
+    no_repeat_ngram_size=2, early_stopping=True
+)
+
+# Decode output
+generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+print(generated_text)
 ```
 
 # Training setup