Update app.py

app.py

@@ -13,17 +13,24 @@ model.to(device)
 
 def historical_generation(prompt, max_new_tokens=600):
     prompt = f"### Text ###\n{prompt}"
-    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    input_ids = inputs["input_ids"].to(device)
+    attention_mask = inputs["attention_mask"].to(device)
 
     # Generate text
-    output = model.generate(input_ids,
-                            max_new_tokens=max_new_tokens,
-                            pad_token_id=tokenizer.eos_token_id,
-                            top_k=50,
-                            temperature=0.3,
-                            top_p=0.95,
-                            do_sample=True,
-                            repetition_penalty=1.5)
+    output = model.generate(
+        input_ids,
+        attention_mask=attention_mask,
+        max_new_tokens=max_new_tokens,
+        pad_token_id=tokenizer.eos_token_id,
+        top_k=50,
+        temperature=0.3,
+        top_p=0.95,
+        do_sample=True,
+        repetition_penalty=1.5,
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id
+    )
 
     # Decode the generated text
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)