Transformers
English
code
Canstralian commited on
Commit
bf27104
·
verified ·
1 Parent(s): 8014c9d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+ import torch
5
+
6
+ # Initialize FastAPI app
7
+ app = FastAPI()
8
+
9
+ # Load the tokenizer and model
10
+ tokenizer = AutoTokenizer.from_pretrained("canstralian/CyberAttackDetection")
11
+ model = AutoModelForCausalLM.from_pretrained("canstralian/CyberAttackDetection")
12
+
13
+ # Define the input data model
14
+ class LogData(BaseModel):
15
+ log: str
16
+
17
+ @app.post("/predict")
18
+ async def predict(data: LogData):
19
+ # Tokenize the input log data
20
+ inputs = tokenizer(data.log, return_tensors="pt")
21
+
22
+ # Generate predictions
23
+ with torch.no_grad():
24
+ outputs = model.generate(**inputs)
25
+
26
+ # Decode the generated tokens to text
27
+ prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
28
+
29
+ return {"prediction": prediction}