Updated backend to handled text-generation parameters

QuoteGenerator.py

@@ -3,7 +3,7 @@ from transformers.pipelines import TextGenerationPipeline
 from typing import Union
 
 class QuoteGenerator():
-    def __init__(self, model_name:str='gruhit13/quote_generator_v1'):
+    def __init__(self, model_name:str='gruhit13/quote-generator-v2'):
         self.model_name = model_name
         self.quote_generator: TextGenerationPipeline
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
@@ -21,14 +21,19 @@ class QuoteGenerator():
         
         return self.tokenizer.bos_token + tags + '<bot>:'
 
-    def generate_quote(self, tags:Union[None, str]=None, 
-                       min_length: int=3, max_length:int=60, 
-                       top_p:float=0.9, top_k:int=5):
+    def generate_quote(self, tags:Union[None, str], max_new_tokens: int, do_sample: bool, 
+                       num_beams: int, top_k: int, top_p: float, temperature: float):
         
         tags = self.preprocess_tags(tags)
-        output = self.quote_generator(tags, min_length=min_length, max_length=max_length,
-                                      temperature=1.0, top_k=5, top_p=top_p, early_stopping=True,
-                                      num_beams=4)
+        output = self.quote_generator(
+            tags, 
+            max_new_tokens=max_new_tokens, 
+            num_beams=num_beams,
+            temperature=temperature, 
+            top_k=top_k, 
+            top_p=top_p, 
+            do_sample = do_sample
+        )
         
         return output[0]['generated_text']
 

main.py

@@ -8,7 +8,6 @@ import os
 
 # API to key to validate the Referer
 API_KEY = os.getenv('API_KEY')
-
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
 # Function to check of the incoming API call is from valid host or not
@@ -20,7 +19,13 @@ def api_key_auth(api_key:str = Depends(oauth2_scheme)):
         )
 
 class QuoteRequest(BaseModel):
-    tags: Union[None, str]
+    tags: Union[None, str] = None
+    do_sample: bool = False
+    max_new_tokens: int = 16
+    num_beams: int = 1
+    top_k: int = 50
+    top_p: float = 1.0
+    temperature: float = 1.0
 
 app = FastAPI()
 
@@ -37,11 +42,19 @@ quote_generator = QuoteGenerator()
 quote_generator.load_generator()
 
 @app.post("/", dependencies=[Depends(api_key_auth)])
-def root(request: Request):
-    return {"message": "This is the website for quote-generator"}
+def root(request: QuoteRequest):
+    print("Incoming request\n", request.__dict__)
+    return {"quote": "<bot>:A beautiful quote generated by bot"}
 
 @app.post("/generate_quote", dependencies=[Depends(api_key_auth)])
 def generate_quote(req: QuoteRequest):
-    print("Tags: ", req.tags)
-    generated_quote_oup = quote_generator.generate_quote(req.tags)
+    generated_quote_oup = quote_generator.generate_quote(
+        tags = req.tags,
+        max_new_tokens = req.max_new_tokens,
+        num_beams = req.num_beams,
+        temperature = req.temperature,
+        top_k = req.top_k,
+        top_p = req.top_p,
+        do_sample = req.do_sample
+    )
     return {'quote': generated_quote_oup}