Hugging Face's logo

Spaces:
dingusagar
/
aita-classifier
Running

App Files Community
dingusagar commited on
Commit
332bb05
·
verified ·
1 Parent(s): 01c41a2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -5,8 +5,9 @@ from ollama import chat
5
  from ollama import ChatResponse
6
 
7
  # Default model
8
- OLLAMA_MODEL = "llama3.2:3b"
9
  # OLLAMA_MODEL = "llama3.2:1b"
 
10
 
11
  # Load BERT MODEL
12
  from transformers import pipeline, DistilBertTokenizerFast
@@ -24,8 +25,8 @@ classifier = pipeline(
24
  )
25
 
26
  bert_label_map = {
27
- 'LABEL_0': 'YTA',
28
- 'LABEL_1': 'NTA',
29
  }
30
 
31
  def ask_bert(prompt):
 
5
  from ollama import ChatResponse
6
 
7
  # Default model
8
+ # OLLAMA_MODEL = "llama3.2:3b"
9
  # OLLAMA_MODEL = "llama3.2:1b"
10
+ OLLAMA_MODEL = "mannix/qwen2-1.5b"
11
 
12
  # Load BERT MODEL
13
  from transformers import pipeline, DistilBertTokenizerFast
 
25
  )
26
 
27
  bert_label_map = {
28
+ 'LABEL_0': 'You are The A-hole (YTA)',
29
+ 'LABEL_1': 'Not the A-hole (NTA)',
30
  }
31
 
32
  def ask_bert(prompt):