File size: 418 Bytes
12c460c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import gradio as gr
from transformers import AutoTokenizer
from model import MultiHeadClassification


model = MultiHeadClassification.from_pretrained('philipp-zettl/multi-head-sequence-classification-model')
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-m3')

def generate(prompt):
    inputs = tokenizer([prompt])
    return model(**inputs)


gr.Interface(
    generate,
    inputs="text",
    ouputs="label"
)