feat: handle multiple prompts

app.py

@@ -15,29 +15,27 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 classifier = pipeline("text-classification", model="saiteki-kai/QA-DeBERTa-v3-large")
 
 @spaces.GPU(duration=60)
-def generate(message):
-    messages = [
-        {"role": "user", "content": message}
-    ]
-    text = tokenizer.apply_chat_template(
+def generate(prompts):
+    messages = [[{"role": "user", "content": message}] for message in prompts]
+
+    texts = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    model_inputs = tokenizer(texts, padding=True, max_new_tokens=512, return_tensors="pt").to(model.device)
     generated_ids = model.generate(
         **model_inputs,
         do_sample=False,
         temperature=0,
         repetition_penalty=1.0,
-        max_new_tokens=512,
     )
     generated_ids = [
         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
     ]
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 
-    return response, classifier(text + "[SEP]" + response)
+    return responses, classifier([text + "[SEP]" + response for text, response in zip(texts, responses)])
 
 
 demo = gr.Interface(fn=generate, inputs=gr.Text(), outputs=gr.Text())