mamiksik commited on
Commit
30d3f5d
1 Parent(s): ec8fdf0
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -1,9 +1,9 @@
1
  import re
2
 
3
  import gradio as gr
4
- from transformers import RobertaTokenizer, pipeline, AutoModelForMaskedLM
5
 
6
- tokenizer = RobertaTokenizer.from_pretrained("mamiksik/CodeBERTa-commit-message-autocomplete")
7
  model = AutoModelForMaskedLM.from_pretrained("mamiksik/CodeBERTa-commit-message-autocomplete")
8
  pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer)
9
 
@@ -47,7 +47,7 @@ def parse_files(patch) -> str:
47
 
48
 
49
  def predict(patch, commit_message):
50
- input_text = parse_files(patch) + "\n<msg> " + commit_message
51
  token_count = tokenizer(input_text, return_tensors="pt").input_ids.shape[1]
52
  result = pipe.predict(input_text)
53
 
 
1
  import re
2
 
3
  import gradio as gr
4
+ from transformers import pipeline, AutoModelForMaskedLM, AutoTokenizer
5
 
6
+ tokenizer = AutoTokenizer.from_pretrained("mamiksik/CodeBERTa-commit-message-autocomplete")
7
  model = AutoModelForMaskedLM.from_pretrained("mamiksik/CodeBERTa-commit-message-autocomplete")
8
  pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer)
9
 
 
47
 
48
 
49
  def predict(patch, commit_message):
50
+ input_text = parse_files(patch) + "</sep></sep><msg> " + commit_message + "</sep>"
51
  token_count = tokenizer(input_text, return_tensors="pt").input_ids.shape[1]
52
  result = pipe.predict(input_text)
53