Took the model definitions outside the functions

app.py

@@ -1,7 +1,13 @@
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 import os
-
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+)
+from optimum.onnxruntime import ORTModelForSeq2SeqLM, ORTModelForSequenceClassification
+from sentence_transformers import SentenceTransformer
 
 # Define the FastAPI app
 app = FastAPI(docs_url="/")
@@ -15,8 +21,32 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Define the Google Books API key
 key = os.environ.get("GOOGLE_BOOKS_API_KEY")
 
+# Define summarization models
+summary_tokenizer_normal = AutoTokenizer.from_pretrained("lidiya/bart-base-samsum")
+summary_model_normal = AutoModelForSeq2SeqLM.from_pretrained("lidiya/bart-base-samsum")
+summary_tokenizer_onnx = AutoTokenizer.from_pretrained("optimum/t5-small")
+summary_model_onnx = ORTModelForSeq2SeqLM.from_pretrained("optimum/t5-small")
+
+# Define classification models
+classification_tokenizer_normal = AutoTokenizer.from_pretrained(
+    "sileod/deberta-v3-base-tasksource-nli"
+)
+classification_model_normal = AutoModelForSequenceClassification.from_pretrained(
+    "sileod/deberta-v3-base-tasksource-nli"
+)
+classification_tokenizer_onnx = AutoTokenizer.from_pretrained(
+    "optimum/distilbert-base-uncased-mnli"
+)
+classification_model_onnx = ORTModelForSequenceClassification.from_pretrained(
+    "optimum/distilbert-base-uncased-mnli"
+)
+
+# Define similarity model
+similarity_model = SentenceTransformer("all-MiniLM-L6-v2")
+
 
 @app.get("/search")
 async def search(
@@ -329,23 +359,14 @@ async def classify(data: list, runtime: str = "normal"):
         pipeline,
     )
     from optimum.onnxruntime import ORTModelForSequenceClassification
-    from optimum.bettertransformer import BetterTransformer
 
     if runtime == "normal":
         # Define the zero-shot classifier
-        tokenizer = AutoTokenizer.from_pretrained(
-            "sileod/deberta-v3-base-tasksource-nli"
-        )
-        model = AutoModelForSequenceClassification.from_pretrained(
-            "sileod/deberta-v3-base-tasksource-nli"
-        )
+        tokenizer = classification_tokenizer_normal
+        model = classification_model_normal
     elif runtime == "onnxruntime":
-        tokenizer = AutoTokenizer.from_pretrained(
-            "optimum/distilbert-base-uncased-mnli"
-        )
-        model = ORTModelForSequenceClassification.from_pretrained(
-            "optimum/distilbert-base-uncased-mnli"
-        )
+        tokenizer = classification_tokenizer_onnx
+        model = classification_model_onnx
 
     classifier_pipe = pipeline(
         "zero-shot-classification",
@@ -396,7 +417,7 @@ async def find_similar(data: list, top_k: int = 5):
         for title, description, publisher in zip(titles, descriptions, publishers)
     ]
 
-    sentence_transformer = SentenceTransformer("all-MiniLM-L6-v2")
+    sentence_transformer = similarity_model
     book_embeddings = sentence_transformer.encode(combined_data, convert_to_tensor=True)
 
     # Make sure that the top_k value is not greater than the number of books
@@ -438,12 +459,12 @@ async def summarize(descriptions: list, runtime="normal"):
 
     # Define the summarizer model and tokenizer
     if runtime == "normal":
-        tokenizer = AutoTokenizer.from_pretrained("lidiya/bart-base-samsum")
-        model = AutoModelForSeq2SeqLM.from_pretrained("lidiya/bart-base-samsum")
+        tokenizer = summary_tokenizer_normal
+        model = summary_model_normal
         model = BetterTransformer.transform(model)
     elif runtime == "onnxruntime":
-        tokenizer = AutoTokenizer.from_pretrained("optimum/t5-small")
-        model = ORTModelForSeq2SeqLM.from_pretrained("optimum/t5-small")
+        tokenizer = summary_tokenizer_onnx
+        model = summary_model_onnx
 
     # Create the summarizer pipeline
     summarizer_pipe = pipeline("summarization", model=model, tokenizer=tokenizer)