| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
| from transformers import BertTokenizer, TFBertModel |
| import keras |
| import traceback |
|
|
| class Prediction(BaseModel): |
| text: str |
|
|
| app = FastAPI() |
| tokenizer = BertTokenizer.from_pretrained('./src/assets/') |
| model = keras.models.load_model('./src/model/scam_class.h5',custom_objects={"TFBertModel": TFBertModel}) |
|
|
| @app.get("/") |
| async def root(): |
| text = "Mainkan terus games mu sekarang dan beli koinnya pakai pulsamu, cek caranya di http://tsel.me/jajanonline" |
| encoded = tokenizer(text=text,add_special_tokens=True,max_length=50,padding='max_length', |
| truncation=True,return_tensors='tf',return_token_type_ids=False,verbose=True,return_attention_mask=True) |
| input_obj = {'input_ids': encoded['input_ids'], 'attention_mask': encoded['attention_mask']} |
| prediction = model.predict(input_obj) |
| pred_arr = prediction.tolist() |
| output = { |
| "neutral": prediction[0][0], |
| "scam": prediction[0][1], |
| "spam": prediction[0][2] |
| } |
| return {"result":pred_arr} |
| @app.post("/predict") |
| async def predict(data: Prediction): |
| try: |
| text = data.text |
| encoded = tokenizer(text=text,add_special_tokens=True,max_length=50,padding='max_length', |
| truncation=True,return_tensors='tf',return_token_type_ids=False,verbose=True,return_attention_mask=True) |
| input_obj = {'input_ids': encoded['input_ids'], 'attention_mask': encoded['attention_mask']} |
| prediction = model.predict(input_obj) |
| pred_arr = prediction.tolist() |
| return {"neutral": pred_arr[0][0], "scam": pred_arr[0][1], "spam": pred_arr[0][2]} |
| except: |
| traceback.print_exc() |
| raise HTTPException(status_code=500, detail="Something went wrong") |
|
|
