File size: 1,160 Bytes
52f03d0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

import torch
import gradio as gr

from typing import Dict
from transformers import pipeline

def food_not_food_classifier(text: str) -> Dict[str, float]:
  food_not_food_classifier = pipeline(task="text-classification",
                                      model="tanvircr7/learn_hf_food_not_food_text_classifier-distilbert-base-uncased",
                                      device="cuda" if torch.cuda.is_available() else "cpu",
                                      top_k=None)
  outputs = food_not_food_classifier(text)[0]
  output_dict = {}
  for item in outputs:
    output_dict[item["label"]] = item["score"]
  return output_dict

description = """
A text classifier to determine if a sentence is about food or not food.
Fine-tuned from DistilBERT
"""

demo = gr.Interface(fn=food_not_food_classifier,
                    inputs="text",
                    outputs=gr.Label(num_top_classes=2),
                    title="Food or Not Food Text Classifier",
                    description=description,
                    examples=[["I whipped up a fresh batch of code"],["A delicious photo of a plate"]]
)

if __name__ == "__main__":
  demo.launch()