code-search-net/code_search_net
Viewer • Updated • 4.14M • 23.8k • 329
How to use UDE-SE/BERTForReturnTypeClassification with Transformers:
# Use a pipeline as a high-level helper
from transformers import pipeline
pipe = pipeline("text-classification", model="UDE-SE/BERTForReturnTypeClassification") # Load model directly
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained("UDE-SE/BERTForReturnTypeClassification")
model = AutoModelForSequenceClassification.from_pretrained("UDE-SE/BERTForReturnTypeClassification")With this model, the return value can be classified based on the method name only. Six classes are available for this purpose. The model answers with the label indexes and the associated probability. The corresponding class name must then be taken from the following list: ['None', 'Number', 'Boolean', 'String', 'Object', 'Collection'].
The corresponding huggingface space can be found at https://huggingface.co/spaces/UDE-SE/ReturnTypePredictor.
Good luck trying it out!