Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 openelm_270m_instruct = AutoModelForCausalLM.from_pretrained("apple/OpenELM-270M-Instruct", trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf")
 
-def generate(newQuestion):
+def generate(newQuestion,num):
     tokenized_prompt = tokenizer(newQuestion)
     tokenized_prompt = torch.tensor(
         tokenized_prompt['input_ids'],
@@ -16,7 +16,7 @@ def generate(newQuestion):
     # Generate
     output_ids = openelm_270m_instruct.generate(
         tokenized_prompt,
-        max_length=128,
+        max_length=int(num),
         pad_token_id=0,
     )
     output_text = tokenizer.decode(
@@ -34,12 +34,12 @@ apple/OpenELM-270M-Instruct
 """
 
 
-def greet(name):
-    return generate(name)
+def greet(name,num):
+    return generate(name,num)
 
 iface = gr.Interface(
     fn=greet,
-    inputs=Textbox(label="Enter Text Here:", type="text"),
+    inputs=[Textbox(label="Enter Text Here:", type="text"),Textbox(label="number of generated tokens:", type="text")],
     outputs=[Textbox(label="generated answer:")],
     title="OpenELM-270M-Instruct",
     # Markdown(developer_info, elem_id="dev-info"),  # Place Markdown directly