Update app.py

app.py

@@ -1,15 +1,12 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
-from load_texts import load_texts
 from tokenizer import SimpleTokenizer
 from transformer import Classifier
 from constants import block_size
 from fastapi.middleware.cors import CORSMiddleware
 
 import uvicorn
-
 import torch
-
 import pickle
 
 app = FastAPI()
@@ -23,35 +20,26 @@ app.add_middleware(
 )
 
 model = None
-tokenizer = None
+tokenizer = SimpleTokenizer()
 pres_dict = {}
 
-with open('pres_dict.pkl', 'rb') as file:
-    reversed_dict = pickle.load(file)
+# load in pres dicts
+with open('speechesdataset/pres_dict.pkl', 'rb') as file1:
+    reversed_dict = pickle.load(file1)
     pres_dict = {value: key for key, value in reversed_dict.items()}
 
-def initialize():
-    global model, tokenizer
-    
-    if not tokenizer:
-        tokenizer = SimpleTokenizer()
-
-        print('start tokenizer')
-        
-        for text in load_texts('train.tsv'):
-            tokenizer.update_vocab(text.split('\t', 1)[1])
-
-        print('finish tokenizer, vocab size is: ', tokenizer.vocab_size)
-        
-    if not model:
-        model = Classifier(tokenizer.vocab_size)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print('loading model')
-        model.load_state_dict(torch.load('all_pres_classifier_model_dict.pth', map_location=device))
-        print('finished loading model')
-        model.to(device)
-        model.eval()
+with open('speechesdataset/tokenizer.pkl', 'rb') as file:
+    tokenizer = pickle.load(file)
 
+# load in model
+model = Classifier(tokenizer.vocab_size)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print('loading model')
+model.load_state_dict(torch.load('speechesdataset/classifier_model_dict.pth', map_location=device))
+print('finished loading model')
+model.to(device)
+model.eval()
+    
 class TextInput(BaseModel):
     text: str
 
@@ -61,11 +49,6 @@ def home():
 
 @app.post("/predict")
 def predict(request: TextInput):
-    global model, tokenizer
-    
-    if model is None or tokenizer is None:
-        initialize()
-    
     text = request.text
 
     # Get the text from the POST request body
@@ -82,5 +65,4 @@ def predict(request: TextInput):
         
     _, predicted = torch.max(output.data, 1)
     
-    return {"predicted": pres_dict[predicted.tolist()[0]]}
-
+    return {"predicted": pres_dict[predicted.tolist()[0]]}