FilesExpand file tree

 master

/

example_model.py

Copy path

BlameMore file actions

BlameMore file actions

Latest commit

 

History

History

History

22 lines (16 loc) · 533 Bytes

 master

/

example_model.py

Top

File metadata and controls

• Code
• Blame

22 lines (16 loc) · 533 Bytes

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

import torch

from transformers import pipeline

from .base_model import BaseModel

class SentimentAnalysisModel(BaseModel):

def __init__(self):

self.model = None

self.load_model()

def load_model(self):

self.model = pipeline("sentiment-analysis")

def predict(self, input_data):

result = self.model(input_data)[0]

return result

if __name__ == "__main__":

sentiment_model = SentimentAnalysisModel()

result = sentiment_model.predict("I love this movie!")

print(result)