Update app.py

app.py

@@ -19,37 +19,38 @@ def chatbot(input):
     
     
     tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
-    
-    # mtp-7b is trained to add "<|endoftext|>" at the end of generations
-    stop_token_ids = tokenizer.convert_tokens_to_ids(["<|endoftext|>"])
-    
-    # define custom stopping criteria object
-    class StopOnTokens(StoppingCriteria):
-        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-            for stop_id in stop_token_ids:
-                if input_ids[0][-1] == stop_id:
-                    return True
-            return False
-    
-    stopping_criteria = StoppingCriteriaList([StopOnTokens()])
-    
-    generate_text = transformers.pipeline(
-        model=model, tokenizer=tokenizer,
-        return_full_text=True,  # langchain expects the full text
-        task='text-generation',
-        device=device,
-        # we pass model parameters here too
-        stopping_criteria=stopping_criteria,  # without this model will ramble
-        temperature=0.1,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max
-        top_p=0.15,  # select from top tokens whose probability add up to 15%
-        top_k=0,  # select from top 0 tokens (because zero, relies on top_p)
-        max_new_tokens=64,  # mex number of tokens to generate in the output
-        repetition_penalty=1.1  # without this output begins repeating
-    )
-    
-    res = generate_text(input)
-    output = res[0]["generated_text"]
-    return output
+
+    for i in range(50):
+        # mtp-7b is trained to add "<|endoftext|>" at the end of generations
+        stop_token_ids = tokenizer.convert_tokens_to_ids(["<|endoftext|>"])
+        
+        # define custom stopping criteria object
+        class StopOnTokens(StoppingCriteria):
+            def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+                for stop_id in stop_token_ids:
+                    if input_ids[0][-1] == stop_id:
+                        return True
+                return False
+        
+        stopping_criteria = StoppingCriteriaList([StopOnTokens()])
+        
+        generate_text = transformers.pipeline(
+            model=model, tokenizer=tokenizer,
+            return_full_text=True,  # langchain expects the full text
+            task='text-generation',
+            device=device,
+            # we pass model parameters here too
+            stopping_criteria=stopping_criteria,  # without this model will ramble
+            temperature=0.1,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max
+            top_p=0.15,  # select from top tokens whose probability add up to 15%
+            top_k=0,  # select from top 0 tokens (because zero, relies on top_p)
+            max_new_tokens=64,  # mex number of tokens to generate in the output
+            repetition_penalty=1.1  # without this output begins repeating
+        )
+        
+        res = generate_text(input)
+        output = res[0]["generated_text"]
+        return output
 
 inputs = gr.inputs.Textbox(lines=7, label="Chat with AI")
 outputs = gr.outputs.Textbox(label="Reply")