Update app.py

app.py

@@ -21,9 +21,21 @@ quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 model = AutoModelForCausalLM.from_pretrained( model_path, quantization_config=quantization_config)
 
 @spaces.GPU
-def generate_text(prompt, temperature, max_length):
-    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(inputs, max_length=max_length, top_p=0.9, temperature=temperature, do_sample=True)
+def generate_text(prompt, temperature=0.9, max_length=1200):
+    # Encode the inputs
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
+    attention_mask = torch.ones(inputs.shape, dtype=torch.long)
+    inputs = inputs.to("cuda")
+    attention_mask = attention_mask.to("cuda")    
+    outputs = model.generate(
+        inputs, 
+        attention_mask=attention_mask,
+        max_length=max_length, 
+        top_p=0.9, 
+        temperature=temperature, 
+        do_sample=True, 
+        pad_token_id=tokenizer.eos_token_id
+    )
     return tokenizer.decode(outputs[0])
 
 def gradio_app():