testing num return sequences

README.md

@@ -5,7 +5,7 @@ colorFrom: gray
 colorTo: purple
 sdk: gradio
 sdk_version: 3.37.0
-app_file: app-ngrok.py
+app_file: app.py
 pinned: true
 license: bigcode-openrail-m
 tags:

app.py

@@ -139,7 +139,7 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
 
     input_ids = tok(messages, return_tensors="pt").input_ids
     input_ids = input_ids.to(m.device)
-    streamer = TextIteratorStreamer(tok, timeout=1000.0, skip_prompt=True, skip_special_tokens=True)
+    #streamer = TextIteratorStreamer(tok, timeout=1000.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         input_ids=input_ids,
         max_new_tokens=max_new_tokens,
@@ -147,27 +147,36 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
         top_p=top_p,
         top_k=top_k,
         repetition_penalty=repetition_penalty,
-        streamer=streamer,
+        #streamer=streamer,
         stopping_criteria=StoppingCriteriaList([stop]),
         num_return_sequences=num_return_sequences,
         num_beams=num_beams,
         do_sample=do_sample,
     )
 
-    stream_complete = Event()
+    #stream_complete = Event()
 
-    def generate_and_signal_complete():
-        m.generate(**generate_kwargs)
-        stream_complete.set()
+    # def generate_and_signal_complete():
+    #     m.generate(**generate_kwargs)
+    #     stream_complete.set()
 
-    t1 = Thread(target=generate_and_signal_complete)
-    t1.start()
+    # t1 = Thread(target=generate_and_signal_complete)
+    # t1.start()
 
-    partial_text = ""
-    for new_text in streamer:
-        partial_text += new_text
+    tokens = m.generate(**generate_kwargs)
 
-    output = format(partial_text) if format_sql else partial_text
+    responses = []
+    for response in tokens:
+        response_text = tok.decode(response, skip_special_tokens=True)
+
+        # Only take what comes after ### Response:
+        response_text = response_text.split("### Response:")[1].strip()
+
+        formatted_text = format(response_text) if format_sql else response_text
+        responses.append(formatted_text)
+
+    # Concat responses to be a single string seperated by a newline
+    output = "\n".join(responses)
 
     if log:
         # Log the request to Firestore