Update app.py

app.py

@@ -5,18 +5,28 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 tokenizer = AutoTokenizer.from_pretrained("TuringsSolutions/Phi3LawCaseManagement", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("TuringsSolutions/Phi3LawCaseManagement", trust_remote_code=True)
 
-def predict(prompt):
+def predict(prompt, temperature, max_tokens):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        temperature=temperature
+    )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
 # Create Gradio interface
-iface = gr.Interface(fn=predict, 
-                     inputs="text", 
-                     outputs="text",
-                     title="Phi3 Law Case Management Model",
-                     description="A model to assist with law case management.")
+iface = gr.Interface(
+    fn=predict, 
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.7, label="Temperature"),
+        gr.inputs.Slider(minimum=10, maximum=200, default=50, step=10, label="Number of Output Tokens")
+    ], 
+    outputs="text",
+    title="Phi3 Law Case Management Model",
+    description="A model to assist with law case management. Adjust the temperature and number of output tokens as needed."
+)
 
 # Launch the Gradio app
-iface.launch()
+iface.launch()