| import pandas as pd | |
| def predict(data, task, model, tokenizer, config, **kwargs): | |
| if isinstance(data, pd.DataFrame): | |
| data = data[data.columns[0]].tolist() | |
| is_df = True | |
| results = [] | |
| addn_args = kwargs.get("addn_args", {}) | |
| for d in data: | |
| inputs = tokenizer(d, return_tensors="pt", return_attention_mask=False) | |
| outputs = model.generate(**inputs, **addn_args, max_length=50) | |
| text = tokenizer.batch_decode(outputs)[0] | |
| results.append(text) | |
| if is_df: | |
| return pd.DataFrame(results,columns =['output']) | |
| return {"output": results} | |