Update src/brain.py

src/brain.py

@@ -6,6 +6,6 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 
 
 def generate_answers(query):
-    input_ids = tokenizer(query, return_tensors="pt")
-    output = model.generate(**input_ids)
-    return tokenizer.decode(output[0])
+    input_ids = tokenizer(query, return_tensors="pt").input_ids
+    output = model.generate(input_ids, max_new_tokens=100)
+    return tokenizer.decode(output[0], skip_special_tokens=True)