Improve inference

app.py

@@ -13,7 +13,7 @@ ffn_dim = embed_dim * 4
 n_vocab = 50257
 max_seq_len = 576
 dropout = 0.1
-top_p_value = 0.6
+top_p_value = 0.8
 
 @spaces.GPU
 def load_model():
@@ -43,7 +43,7 @@ app = gr.Interface(
     [
         gr.Textbox(label='Prompt', lines=3), 
         gr.Number(label='Output Tokens', value=100),
-        gr.Slider(0.1, 2.0, step=0.1, value=0.7, label='Temperature')
+        gr.Slider(0.1, 2.0, step=0.1, value=0.8, label='Temperature')
     ],
     gr.Textbox(label='Output', lines=11)
 )

model.py

@@ -60,8 +60,8 @@ class GPTModel(nn.Module):
                 logits = self(context)[0,-1,:]
 
                 if top_p:
-                    probs = torch.nn.functional.softmax(logits, dim=-1)
-                    sorted_probs, sorted_indices = torch.sort(probs, dim=-1)
+                    probs = torch.nn.functional.softmax(logits / temperature, dim=-1)
+                    sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
                     cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
                     sorted_mask = cumulative_probs <= top_p
                     sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
@@ -71,8 +71,9 @@ class GPTModel(nn.Module):
                     probs_sampled = torch.multinomial(filtered_probs, 1).item()
                     selected_token = sorted_indices[probs_sampled].item()
                 elif top_k:
-                    topk_probs, topk_indices = logits.topk(topk_elements)
-                    probs = nn.functional.softmax(topk_probs / temperature, dim=-1)
+                    scaled_logits = logits / temperature
+                    topk_probs, topk_indices = scaled_logits.topk(topk_elements)
+                    probs = nn.functional.softmax(topk_probs, dim=-1)
                     probs_sampled = torch.multinomial(probs, 1).item()
                     selected_token = topk_indices[probs_sampled].item()
                 else: # Greedy decoding