Changed some loops to list comprehension

app.py

@@ -1,8 +1,19 @@
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 
 # Define the FastAPI app
 app = FastAPI(docs_url="/")
 
+# Add the CORS middleware to the app
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
 @app.get("/search={query}&similarity={similarity}")
 def search(query, similarity="false"):
     import time
@@ -12,7 +23,7 @@ def search(query, similarity="false"):
 
     # Set the API endpoint and query parameters
     url = "https://www.googleapis.com/books/v1/volumes"
-    params = {"q": str(query), "printType": "books", "maxResults": 30}
+    params = {"q": str(query), "printType": "books", "maxResults": 10}
 
     # Send a GET request to the API with the specified parameters
     response = requests.get(url, params=params)
@@ -141,7 +152,7 @@ def search(query, similarity="false"):
         try:
             # Parse the result
             parsed_result = parse_result(result, ordered_keys=ordered_keys)
- 
+
             # Append the parsed result to the lists
             titles.append(parsed_result["Title"])
             authors.append(parsed_result["Author"])
@@ -155,7 +166,6 @@ def search(query, similarity="false"):
         except IndexError:
             break
 
-
     ### Prediction ###
     from transformers import (
         AutoTokenizer,
@@ -163,8 +173,7 @@ def search(query, similarity="false"):
         AutoModelForSequenceClassification,
         pipeline,
     )
-    from sentence_transformers import SentenceTransformer, CrossEncoder
-    from sentence_transformers.util import cos_sim, dot_score
+    from sentence_transformers import SentenceTransformer
 
     # Load the classifiers
     # classifier = TextClassifier.load(
@@ -175,7 +184,7 @@ def search(query, similarity="false"):
 
     # Combine title, description, and publisher into a single string
     combined_data = [
-        f"The book's title is {title.title()}. It is published by {publisher.title()}. This book is about {description}"
+        f"The book's title is {title}. It is published by {publisher}. This book is about {description}"
         for title, description, publisher in zip(titles, descriptions, publishers)
     ]
 
@@ -191,10 +200,10 @@ def search(query, similarity="false"):
     # classes = [sentence.labels for sentence in sentences]
 
     # Define the summarizer model and tokenizer
-    sum_tokenizer = AutoTokenizer.from_pretrained("lidiya/bart-base-samsum")
+    sum_tokenizer = AutoTokenizer.from_pretrained("sshleifer/distilbart-xsum-12-6")
 
-    # sum_model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-xsum-12-6")
-    sum_model = AutoModelForSeq2SeqLM.from_pretrained("lidiya/bart-base-samsum")
+    sum_model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distilbart-xsum-12-6")
+    # sum_model = AutoModelForSeq2SeqLM.from_pretrained("lidiya/bart-base-samsum")
 
     summarizer_pipeline = pipeline(
         "summarization",
@@ -240,7 +249,7 @@ def search(query, similarity="false"):
     ]
 
     # Get the predicted labels
-    classes = zs_classifier(combined_data, candidate_labels)
+    classes = [zs_classifier(doc, candidate_labels) for doc in combined_data]
 
     # Calculate the elapsed time
     end_time = time.time()
@@ -272,21 +281,20 @@ def search(query, similarity="false"):
         similar_books = [{"sorted_by_similarity": []} for i in range(len(titles))]
 
     # Create a list of dictionaries to store the results
-    results = []
-    for i in range(len(titles)):
-        results.append(
-            {
-                "id": i,
-                "title": titles[i],
-                "author": authors[i],
-                "publisher": publishers[i],
-                "image_link": images[i],
-                "labels": classes[i]["labels"][0:2],
-                "label_confidences": classes[i]["scores"][0:2],
-                "summary": summaries[i][0]["summary_text"],
-                "similar_books": similar_books[i]["sorted_by_similarity"],
-                "runtime": runtime,
-            }
-        )
+    results = [
+        {
+            "id": i,
+            "title": titles[i],
+            "author": authors[i],
+            "publisher": publishers[i],
+            "image_link": images[i],
+            "labels": classes[i]["labels"][0:2],
+            "label_confidences": classes[i]["scores"][0:2],
+            "summary": summaries[i][0]["summary_text"],
+            "similar_books": similar_books[i]["sorted_by_similarity"],
+            "runtime": runtime,
+        }
+        for i in range(len(titles))
+    ]
 
     return results