Update app.py

app.py

@@ -14,7 +14,7 @@ model.eval()
 model.to('cpu')
 
 # Define the function that generates text from a prompt
-def generate_text(prompt):
+def generate_text(prompt, temperature, top_p):
     input_tokens = tokenizer.encode(prompt, return_tensors='pt')
     input_tokens = input_tokens.to('cpu')
     
@@ -23,15 +23,34 @@ def generate_text(prompt):
     for _ in range(80):  # Adjust the range to control the number of tokens generated
         with torch.no_grad():
             outputs = model(input_tokens)
-            predictions = outputs.logits
+            predictions = outputs.logits / temperature
+            sorted_logits, sorted_indices = torch.sort(predictions[:, -1, :], descending=True)
+            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+            sorted_indices_to_remove = cumulative_probs > top_p
+            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+            sorted_indices_to_remove[..., 0] = 0
+            indices_to_remove = sorted_indices[sorted_indices_to_remove]
+            predictions[:, -1, indices_to_remove] = -float('Inf')
             next_token = torch.multinomial(torch.softmax(predictions[:, -1, :], dim=-1), 1)
         
         input_tokens = torch.cat((input_tokens, next_token), dim=1)
         
         decoded_token = tokenizer.decode(next_token.item())
         generated_text += decoded_token  # Append the new token to the generated text
+        if decoded_token == "#":  # Stop if the end of sequence token is generated
+            break
         yield generated_text  # Yield the entire generated text so far
 
-# Create a Gradio interface with a text input and a text output
-interface = gr.Interface(fn=generate_text, inputs='text', outputs='text', live=False)
-interface.launch()
+# Create a Gradio interface with a text input, sliders for temperature and top_p, and a text output
+interface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=1.0, label="Temperature"),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.9, label="Top-p")
+    ],
+    outputs='text',
+    live=False
+)
+
+interface.launch()