Update handler.py

handler.py

@@ -94,12 +94,14 @@ class EndpointHandler():
 
         max_new_tokens=5
         do_sample=False
-        topk=3
+        topk=5
         base_lambda=1.
+        
         inputs = query.pop("inputs", query)
-        logger.info(type(inputs))
-        if isinstance(inputs, str):
-            inputs = [inputs]
+        # if isinstance(inputs, str):
+            # inputs = [inputs]
+        if not inputs.endswith("<cite>"):
+            inputs += "<cite>"
         inputs = self.tokenizer(inputs, return_tensors="pt").to("cuda")
         with torch.no_grad():
             outputs = self.model(**inputs)