Update app.py

app.py

@@ -36,39 +36,53 @@ model = load_model()
 safety = SafetyManager(model, tokenizer)
 max_response_length = 200
 
-def generate_response(model, tokenizer, input_ids, safety_manager, max_length=200):
+def generate_response(model, tokenizer, input_ids, safety_manager, max_length=200, temperature=1.0):
     device = next(model.parameters()).device
     generated_ids = input_ids.copy()
     for _ in range(max_length):
         input_tensor = torch.tensor([generated_ids], device=device)
         with torch.no_grad():
             logits = model(input_tensor)
-        next_token = logits.argmax(-1)[:, -1].item()
+        
+        # Get last token logits and apply temperature
+        next_token_logits = logits[0, -1, :]
+        if temperature != 1.0:
+            next_token_logits = next_token_logits / temperature
+        probs = torch.softmax(next_token_logits, dim=-1)
+        
+        # Sample next token
+        next_token = torch.multinomial(probs, num_samples=1).item()
+        
+        # Stop if end token is generated
         if next_token == tokenizer.token_to_id("</s>"):
             break
+            
+        # Safety check
         current_text = tokenizer.decode(generated_ids + [next_token])
         if not safety_manager.content_filter(current_text):
             break
+            
         generated_ids.append(next_token)
     return generated_ids[len(input_ids):]
 
-def process_message(user_input, chat_history, token_history):
+def process_message(user_input, chat_history, token_history, temperature, max_context_length):
     # Process user input
     user_turn = f"<user> {user_input} </s>"
     user_tokens = tokenizer.encode(user_turn).ids
     token_history.extend(user_tokens)
     
-    # Prepare input sequence
+    # Prepare input sequence with context limit
     input_sequence = [tokenizer.token_to_id("<s>")] + token_history
     
-    # Truncate if needed
-    max_input_len = CONFIG["max_seq_len"] - max_response_length
+    # Truncate based on max context length
+    max_input_len = max_context_length
     if len(input_sequence) > max_input_len:
         input_sequence = input_sequence[-max_input_len:]
         token_history = input_sequence[1:]
     
-    # Generate response
-    response_ids = generate_response(model, tokenizer, input_sequence, safety, max_response_length)
+    # Generate response with temperature
+    response_ids = generate_response(model, tokenizer, input_sequence, safety, 
+                                    max_response_length, temperature)
     
     # Process assistant response
     assistant_text = "I couldn't generate a proper response."
@@ -97,9 +111,16 @@ with gr.Blocks() as demo:
     msg = gr.Textbox(label="Your Message")
     token_state = gr.State([])
     
+    with gr.Row():
+        temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, 
+                              label="Temperature (higher = more random)")
+        max_context = gr.Slider(100, CONFIG["max_seq_len"] - max_response_length, 
+                              value=CONFIG["max_seq_len"] - max_response_length, step=1,
+                              label="Max Context Length")
+    
     msg.submit(
         process_message,
-        [msg, chatbot, token_state],
+        [msg, chatbot, token_state, temperature, max_context],
         [chatbot, token_state],
         queue=False
     ).then(