Pujan-Dev commited on
Commit
b59d3a6
·
verified ·
1 Parent(s): 7dbfd9c

updated app.py

Browse files
Files changed (1) hide show
  1. app.py +18 -8
app.py CHANGED
@@ -1,9 +1,10 @@
1
- import torch
2
- from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
3
- from fastapi import FastAPI, HTTPException
4
  from pydantic import BaseModel
5
- from contextlib import asynccontextmanager
 
6
  import asyncio
 
7
 
8
  # FastAPI app instance
9
  app = FastAPI()
@@ -11,6 +12,9 @@ app = FastAPI()
11
  # Global model and tokenizer variables
12
  model, tokenizer = None, None
13
 
 
 
 
14
  # Function to load model and tokenizer
15
  def load_model():
16
  model_path = "./Ai-Text-Detector/model"
@@ -21,7 +25,7 @@ def load_model():
21
  config = GPT2Config.from_pretrained(model_path)
22
  model = GPT2LMHeadModel(config)
23
  model.load_state_dict(torch.load(weights_path, map_location=torch.device("cpu")))
24
- model.eval() # Set model to evaluation mode
25
  except Exception as e:
26
  raise RuntimeError(f"Error loading model: {str(e)}")
27
 
@@ -61,14 +65,20 @@ def classify_text(sentence: str):
61
 
62
  return result, perplexity
63
 
64
- # POST route to analyze text
65
  @app.post("/analyze")
66
- async def analyze_text(data: TextInput):
67
  user_input = data.text.strip()
 
68
  if not user_input:
69
  raise HTTPException(status_code=400, detail="Text cannot be empty")
 
 
 
 
 
 
70
 
71
- # Run classification asynchronously to prevent blocking
72
  result, perplexity = await asyncio.to_thread(classify_text, user_input)
73
 
74
  return {
 
1
+ from fastapi import FastAPI, HTTPException, Depends
2
+ from fastapi.security import HTTPBearer
 
3
  from pydantic import BaseModel
4
+ from transformers import GPT2LMHeadModel, GPT2TokenizerFast, GPT2Config
5
+ import torch
6
  import asyncio
7
+ from contextlib import asynccontextmanager
8
 
9
  # FastAPI app instance
10
  app = FastAPI()
 
12
  # Global model and tokenizer variables
13
  model, tokenizer = None, None
14
 
15
+ # HTTPBearer instance for security
16
+ bearer_scheme = HTTPBearer()
17
+
18
  # Function to load model and tokenizer
19
  def load_model():
20
  model_path = "./Ai-Text-Detector/model"
 
25
  config = GPT2Config.from_pretrained(model_path)
26
  model = GPT2LMHeadModel(config)
27
  model.load_state_dict(torch.load(weights_path, map_location=torch.device("cpu")))
28
+ model.eval()
29
  except Exception as e:
30
  raise RuntimeError(f"Error loading model: {str(e)}")
31
 
 
65
 
66
  return result, perplexity
67
 
68
+ # POST route to analyze text with Bearer token
69
  @app.post("/analyze")
70
+ async def analyze_text(data: TextInput, token: str = Depends(bearer_scheme)):
71
  user_input = data.text.strip()
72
+
73
  if not user_input:
74
  raise HTTPException(status_code=400, detail="Text cannot be empty")
75
+
76
+ # Check if there are at least two words
77
+ word_count = len(user_input.split())
78
+ if word_count < 2:
79
+ raise HTTPException(status_code=400, detail="Text must contain at least two words")
80
+
81
 
 
82
  result, perplexity = await asyncio.to_thread(classify_text, user_input)
83
 
84
  return {