Update app.py

app.py

@@ -1,40 +1,52 @@
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoTokenizer, AutoModel
 
-# Load the local model
+# Load tokenizer and model
 model_name = "codewithdark/latent-recurrent-depth-lm"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device).eval()  # Set to evaluation mode
+model.to(device).eval()
 
-# Define inference function
-def chat_with_model(input_text, model_choice):
-    if model_choice == "Latent Recurrent Depth LM":
-        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-        with torch.no_grad():
-            output = model.generate(input_ids, max_length=512)
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        return response
-    return "Model not available yet!"
+# Define function for inference
+def chat_with_model(input_text, num_iterations, max_tokens, temperature, top_k):
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
 
-# Create Gradio Interface
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids,
+            max_length=max_tokens,
+            num_iterations=num_iterations,  # Assuming the model supports it
+            temperature=temperature,
+            top_k=top_k
+        )
+    
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return response
+
+# Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 Chat with Latent Recurrent Depth LM")
+
+    with gr.Row():
+        text_input = gr.Textbox(label="Enter your message")
     
-    model_choice = gr.Radio(
-        ["Latent Recurrent Depth LM"],  # Add more models if needed
-        label="Select Model",
-        value="Latent Recurrent Depth LM"
-    )
+    with gr.Row():
+        num_iterations = gr.Slider(1, 20, step=1, value=10, label="Number of Iterations")
+        max_tokens = gr.Slider(10, 200, step=10, value=50, label="Max Tokens")
+        temperature = gr.Slider(0.1, 1.0, step=0.1, value=0.5, label="Temperature")
+        top_k = gr.Slider(10, 100, step=10, value=50, label="Top-K Sampling")
     
-    text_input = gr.Textbox(label="Enter your message")
     submit_button = gr.Button("Generate Response")
     output_text = gr.Textbox(label="Model Response")
 
-    submit_button.click(fn=chat_with_model, inputs=[text_input, model_choice], outputs=output_text)
+    submit_button.click(
+        fn=chat_with_model,
+        inputs=[text_input, num_iterations, max_tokens, temperature, top_k],
+        outputs=output_text
+    )
 
-# Launch the Gradio app
+# Launch Gradio app
 if __name__ == "__main__":
     demo.launch()