nsethi610 commited on
Commit
b8dad8b
1 Parent(s): f0a38b0

Create task.py

Browse files
Files changed (1) hide show
  1. task.py +34 -0
task.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ class Task:
3
+ def __init__(self, label, name) -> None:
4
+ self.name = name
5
+ self.label = label
6
+ self.models = []
7
+ self.examples = []
8
+
9
+ def initialize(self, models, examples):
10
+ self.models = models
11
+ self.examples = examples
12
+
13
+ def __str__(self) -> str:
14
+ return f"Name: {self.name}, Label: {self.label}, Models: {self.models}, examples: {self.examples}"
15
+
16
+
17
+ tasks_config = {
18
+ "text-generation": {"name": "Text Generation", "config": {
19
+ "models": ["meta-llama/Llama-2-7b-chat-hf", "mistralai/Mixtral-8x7B-Instruct-v0.1"]
20
+ }, "info": " The idea here is that you provide a prompt and the model will auto-complete it by generating the remaining text. This is similar to the predictive text feature that is found on many phones. Text generation involves randomness, so it’s normal if you don’t get the desired results."
21
+ },
22
+ "fill-mask": {"name": "Fill Mask", "config": {
23
+ "models": ["google-bert/bert-base-uncased", "distilbert/distilbert-base-uncased"]
24
+ }, "info": "The idea of this task is to fill in the blanks in a given text. Example: This course will teach you all about <mask> models. and pipeline would try to predict <mask>."},
25
+ "summarization": {"name": "Summarization", "config": {
26
+ "models": ["facebook/bart-large-cnn", "Falconsai/text_summarization"]
27
+ }, "info": "Summarization is the task of reducing a text into a shorter text while keeping all (or most) of the important aspects referenced in the text"},
28
+ "ner": {"name": "Named Entity Recognition", "config": {
29
+ "models": ["dslim/bert-base-NER", "Jean-Baptiste/roberta-large-ner-english"]
30
+ }, "info": "Named entity recognition (ner) is a task where the model has to find which parts of the input text correspond to entities such as persons, locations, or organizations."},
31
+ "question-answering": {"name": "Question Answering", "config": {
32
+ "models": ["deepset/roberta-base-squad2", "timpal0l/mdeberta-v3-base-squad2"]
33
+ }, "info": "The question-answering pipeline answers questions using information from a given context"}
34
+ }