Update app.py

app.py

@@ -15,11 +15,18 @@ tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_token)
 model = AutoModelForCausalLM.from_pretrained(model_name, token=api_token)
 
 # Define the inference function
-def generate_text(prompt, max_length=100, temperature=0.7):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    output = model.generate(inputs['input_ids'], max_length=max_length, temperature=temperature)
+def generate_text(prompt, max_length, temperature):
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    output = model.generate(
+        inputs['input_ids'],
+        max_length=max_length,
+        temperature=temperature,
+        do_sample=True,  # Enable sampling
+        attention_mask=inputs['attention_mask']  # Pass attention mask
+    )
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
+
 # Create the Gradio interface
 iface = gr.Interface(
     fn=generate_text,