File size: 599 Bytes
503dc31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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}