Choose 1st response

app.py

@@ -144,6 +144,7 @@ def generate_dummy_db(db_info, question):
             response_text = response['choices'][0]['message']['content']
             
             db_code = extract_db_code(response_text)
+            print(db_code)
 
             return db_code
             
@@ -183,6 +184,9 @@ def test_query_on_dummy_db(db_code, query):
 
 
 def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, log=False, num_return_sequences=1, num_beams=1, do_sample=False):
+    if num_return_sequences >= num_beams:
+        gr.Warning("Num return sequences must be less than or equal to num beams.")
+    
     stop_token_ids = tok.convert_tokens_to_ids(["###"])
     class StopOnTokens(StoppingCriteria):
         def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
@@ -237,7 +241,7 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
             responses.append(query)
             
     # Choose a random response from responses
-    output = responses[random.randint(0, len(responses)-1)] if len(responses) > 0 else "###"
+    output = responses[0] if len(responses) > 0 else "###"
 
     if log:
         # Log the request to Firestore