added main and test codes

main.py

@@ -1,6 +1,126 @@
-def main():
-    print("Hello from sentiment-classification!")
+import uuid
+import torch
+import asyncio
+import transformers
+from typing import Dict
+from fastapi import FastAPI
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+from transformers import (
+    pipeline,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    BitsAndBytesConfig,
+)
 
+# ----------------------------- #
+#         Configurations        #
+# ----------------------------- #
+transformers.set_seed(42)
+torch.set_default_dtype(torch.bfloat16)
 
-if __name__ == "__main__":
-    main()
+MODEL_NAME = "climatebert/distilroberta-base-climate-sentiment"
+BATCH_PROCESS_INTERVAL = 0.01
+MAX_BATCH_SIZE = 128
+
+# ----------------------------- #
+#        Shared Storage         #
+# ----------------------------- #
+query_queue: asyncio.Queue = asyncio.Queue()
+results: Dict[str, Dict] = {}
+classifier = None  # will be initialized in lifespan
+
+
+# ----------------------------- #
+#      Model Initialization     #
+# ----------------------------- #
+def load_classifier(model_name: str):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        device_map="auto",
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        ),
+    )
+    return pipeline(
+        "text-classification", model=model, tokenizer=tokenizer, framework="pt"
+    )
+
+
+# ----------------------------- #
+#         Pydantic Schema       #
+# ----------------------------- #
+class Query(BaseModel):
+    sentence: str
+
+
+# ----------------------------- #
+#      Queue Processing Task    #
+# ----------------------------- #
+async def process_queue():
+    while True:
+        await asyncio.sleep(BATCH_PROCESS_INTERVAL)
+
+        batch = []
+        while not query_queue.empty() and len(batch) < MAX_BATCH_SIZE:
+            batch.append(await query_queue.get())
+
+        if not batch:
+            continue
+
+        sentences = [item["sentence"] for item in batch]
+        ids = [item["id"] for item in batch]
+        predictions = classifier(sentences, batch_size=len(sentences))
+
+        for query_id, pred, sentence in zip(ids, predictions, sentences):
+            results[query_id] = {
+                "sentence": sentence,
+                "label": pred["label"],
+                "score": pred["score"],
+            }
+
+
+# ----------------------------- #
+#        Lifespan Handler       #
+# ----------------------------- #
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global classifier
+    classifier = load_classifier(MODEL_NAME)
+    _ = classifier("Startup warm-up sentence.")
+    queue_task = asyncio.create_task(process_queue())
+    yield
+    queue_task.cancel()
+    try:
+        await queue_task
+    except asyncio.CancelledError:
+        pass
+
+
+# ----------------------------- #
+#         FastAPI Setup         #
+# ----------------------------- #
+app = FastAPI(lifespan=lifespan)
+
+
+# ----------------------------- #
+#         API Endpoints         #
+# ----------------------------- #
+@app.post("/classify")
+async def classify(query: Query):
+    query_id = str(uuid.uuid4())
+    await query_queue.put({"id": query_id, "sentence": query.sentence})
+
+    while query_id not in results:
+        await asyncio.sleep(0.001)
+
+    return {"id": query_id, "result": results.pop(query_id)}
+
+
+@app.get("/")
+def read_root():
+    return {
+        "message": "Welcome to the Sentiment Classification API with Query Batching"
+    }

test.py

@@ -0,0 +1,99 @@
+import os
+import time
+import random
+import requests
+import logging
+from typing import Union, Tuple
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm  # optional: install via `pip install tqdm`
+
+# ----------------------------- #
+#        Configuration          #
+# ----------------------------- #
+URL = "http://localhost:8000/classify"
+NUM_REQUESTS = 4000
+MAX_WORKERS = os.cpu_count() * 8 or 2
+TIMEOUT = 20
+
+logging.basicConfig(level=logging.INFO)
+
+# Sample text tokens
+SAMPLE_POPULATION = """Limes have higher contents of sugars and acids than lemons do.[1] Lime juice may be squeezed from fresh limes, or purchased in bottles in both unsweetened and sweetened varieties. Lime juice is used to make limeade, and as an ingredient (typically as sour mix) in many cocktails.
+
+Lime pickles are an integral part of Indian cuisine, especially in South India. In Kerala, the Onam Sadhya usually includes either lemon pickle or lime pickle. Other Indian preparations of limes include sweetened lime pickle, salted pickle, and lime chutney.
+
+In cooking, lime is valued both for the acidity of its juice and the floral aroma of its zest. It is a common ingredient in authentic Mexican, Vietnamese and Thai dishes. Lime soup is a traditional dish from the Mexican state of Yucatan. It is also used for its pickling properties in ceviche. Some guacamole recipes call for lime juice.
+
+The use of dried limes (called black lime or limoo) as a flavouring is typical of Persian cuisine, Iraqi cuisine, as well as in Eastern Arabian cuisine baharat (a spice mixture that is also called kabsa or kebsa).
+
+Key lime gives the character flavouring to the American dessert known as Key lime pie. In Australia, desert lime is used for making marmalade.
+
+Lime is an ingredient in several highball cocktails, often based on gin, such as gin and tonic, the gimlet and the Rickey. Freshly squeezed lime juice is also considered a key ingredient in margaritas, although sometimes lemon juice is substituted. It is also found in many rum cocktails such as the daiquiri, and other tropical drinks.
+
+Lime extracts and lime essential oils are frequently used in perfumes, cleaning products, and aromatherapy.""".split()
+
+
+# ----------------------------- #
+#        Request Builder        #
+# ----------------------------- #
+def build_payload() -> dict:
+    sentence = " ".join(
+        random.choices(SAMPLE_POPULATION, k=random.randint(20, len(SAMPLE_POPULATION)))
+    )
+    return {"sentence": sentence}
+
+
+# ----------------------------- #
+#      Request Sender Logic     #
+# ----------------------------- #
+def send_request() -> Union[int, str]:
+    try:
+        response = requests.post(
+            URL,
+            json=build_payload(),
+            headers={"Content-Type": "application/json"},
+            timeout=TIMEOUT,
+        )
+        return response.status_code
+    except requests.RequestException as e:
+        return f"Error: {e}"
+
+
+# ----------------------------- #
+#         Test Runner           #
+# ----------------------------- #
+def test_endpoint():
+    print(f"Sending {NUM_REQUESTS} requests to {URL} with {MAX_WORKERS} workers")
+    start_time = time.time()
+
+    successful = 0
+    failed = 0
+    status_distribution = {}
+
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        futures = [executor.submit(send_request) for _ in range(NUM_REQUESTS)]
+        for future in tqdm(
+            as_completed(futures), total=NUM_REQUESTS, desc="Processing"
+        ):
+            result = future.result()
+            if isinstance(result, int):
+                status_distribution[result] = status_distribution.get(result, 0) + 1
+                if 200 <= result < 300:
+                    successful += 1
+                else:
+                    failed += 1
+            else:
+                failed += 1
+                logging.warning(result)
+
+    duration = time.time() - start_time
+    print("\n--- Test Summary ---")
+    print(f"Elapsed Time        : {duration:.2f} seconds")
+    print(f"Total Requests Sent : {NUM_REQUESTS}")
+    print(f"Successful Requests : {successful}")
+    print(f"Failed Requests     : {failed}")
+    print(f"Status Code Summary : {status_distribution}")
+
+
+if __name__ == "__main__":
+    test_endpoint()