camilin29/github_pull_request_classifier

1from fastapi import FastAPI2from pydantic import BaseModel3from transformers import RobertaTokenizer4from transformers import AutoConfig, AutoModel5from torch.utils.data import DataLoader6import torch7 8# define device9device = 'cuda' if torch.cuda.is_available() else 'cpu'10 11 12# Load the model and tokenizer13 14class BERTClass(torch.nn.Module):15    def __init__(self):16        super(BERTClass, self).__init__()17        self.config = AutoConfig.from_pretrained('roberta-base')18        self.bert_model = AutoModel.from_pretrained('roberta-base', return_dict=True)19        self.dropout = torch.nn.Dropout(0.3)20        self.linear = torch.nn.Linear(768, 4)21 22    def forward(self, ids, mask, token_type_ids):23        output = self.bert_model(24            ids,25            attention_mask=mask,26            token_type_ids=token_type_ids27        )28 29        output_dropout = self.dropout(output.pooler_output)30        output = self.linear(output_dropout)31        return output32 33 34model = BERTClass()35model.load_state_dict(torch.load('roberta_model.pth'))36model.to('cuda')37 38tokenizer = RobertaTokenizer.from_pretrained('roberta_tokenizer', local_files_only=True)39 40app = FastAPI()41 42 43@app.get("/predict")44async def predict(input_text: str):45    # Encode the input text46    input_ids = tokenizer.encode(str(input_text), return_tensors="pt")47 48    # Add attention mask49    attention_mask = input_ids.ne(tokenizer.pad_token_id)50 51    # Set token type ids to zeros (for a single sentence)52    token_type_ids = torch.zeros_like(input_ids)53 54    # Pass the input through the model55    output = model(input_ids.to(device), attention_mask.to(device), token_type_ids.to(device))56 57    # Get the prediction58    prediction = torch.argmax(output[0]).item()59 60    dict_values = {61        0: 'deprecated',62        1: 'features',63        2: 'fix',64        3: 'maintenance'}65    return {"prediction": dict_values.get(prediction)}66