YOGEXH commited on
Commit
486f81d
·
verified ·
1 Parent(s): fb32211

Upload 5 files

Browse files
api.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import AlbertForSequenceClassification, AlbertTokenizer
4
+ import torch
5
+
6
+ app = FastAPI()
7
+
8
+ MODEL_PATH = "./saved_model"
9
+
10
+ tokenizer = AlbertTokenizer.from_pretrained(MODEL_PATH)
11
+ model = AlbertForSequenceClassification.from_pretrained(MODEL_PATH)
12
+
13
+ model.eval()
14
+
15
+ class JobInput(BaseModel):
16
+ description: str
17
+
18
+ @app.get("/")
19
+ def read_root():
20
+ return {
21
+ "message": "ScamZero Prediction API is running!",
22
+ "endpoints": {
23
+ "predict": "/predict (POST)",
24
+ "health": "/health (GET)"
25
+ }
26
+ }
27
+
28
+ @app.get("/health")
29
+ def health_check():
30
+ return {"status": "healthy", "model": "albert-base-v2"}
31
+
32
+ @app.post("/predict")
33
+ def predict(data: JobInput):
34
+
35
+ text = data.description
36
+
37
+ inputs = tokenizer(
38
+ text,
39
+ return_tensors="pt",
40
+ padding=True,
41
+ truncation=True,
42
+ max_length=256
43
+ )
44
+
45
+ with torch.no_grad():
46
+ outputs = model(**inputs)
47
+
48
+ probs = torch.softmax(outputs.logits, dim=-1)
49
+ prediction = torch.argmax(probs, dim=-1).item()
50
+ confidence = probs[0][prediction].item()
51
+
52
+ return {
53
+ "prediction": prediction,
54
+ "label": "Scam" if prediction == 1 else "Safe",
55
+ "confidence": confidence
56
+ }
saved_model/config.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0b8d9198c9d665b8002c4e85da3284893bfa79d8fbef728bcf5617ae3c23d6d5
3
+ size 981
saved_model/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:90bcca679551ab1bfb5b156332b03a2776bc2d01769ba6f15021eb9188d426ff
3
+ size 46743904
saved_model/tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9870e912be4f084ce4a3eb0db23b330446da26b97aecc77a99f4f79c29e3cc39
3
+ size 1957992
saved_model/tokenizer_config.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:be010b6240b82e80033768290ec1eba4c98a577b91647ed45e14e8262ead5c29
3
+ size 407