Spaces:
Running
Running
| from fastapi import FastAPI | |
| from pydantic import BaseModel | |
| from transformers import BertTokenizer, BertForSequenceClassification | |
| import torch | |
| import uvicorn # <--- THIS IS THE FIX | |
| # --- Configuration --- | |
| MODEL_PATH = "./BERT_Bullying_Detector_Model" | |
| device = torch.device("cpu") | |
| # --- Data Model for Input --- | |
| class TextInput(BaseModel): | |
| text: str | |
| # --- Load Model and Tokenizer --- | |
| model = None | |
| tokenizer = None | |
| try: | |
| model = BertForSequenceClassification.from_pretrained(MODEL_PATH) | |
| tokenizer = BertTokenizer.from_pretrained(MODEL_PATH) | |
| model.to(device) | |
| model.eval() | |
| print("Model and tokenizer loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading model: {e}") | |
| # --- Create FastAPI App --- | |
| app = FastAPI() | |
| def read_root(): | |
| return {"status": "API is running"} | |
| def predict_toxicity(input_data: TextInput): | |
| if not model or not tokenizer: | |
| return {"error": "Model not loaded."} | |
| text = input_data.text | |
| try: | |
| encoding = tokenizer.encode_plus( | |
| text, | |
| add_special_tokens=True, | |
| max_length=128, | |
| return_token_type_ids=False, | |
| padding='max_length', | |
| return_tensors='pt', | |
| truncation=True | |
| ) | |
| input_ids = encoding['input_ids'].to(device) | |
| attention_mask = encoding['attention_mask'].to(device) | |
| with torch.no_grad(): | |
| outputs = model(input_ids=input_ids, attention_mask=attention_mask) | |
| _, prediction = torch.max(outputs.logits, dim=1) | |
| label_map = {0: "Not Bullying", 1: "Bullying"} | |
| result_label = label_map[prediction.item()] | |
| return {"label": result_label, "score": prediction.item()} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| # --- Start the Server --- | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |