KhairulAmirinUM commited on
Commit
223c527
·
1 Parent(s): 0ffbeb2
Files changed (1) hide show
  1. src/hf.py +32 -32
src/hf.py CHANGED
@@ -2,41 +2,41 @@ import os.path
2
 
3
  from transformers import BertTokenizer, BertForSequenceClassification,TextClassificationPipeline, AutoModelForSequenceClassification
4
  # Load tokenizer and model from the fine-tuned directory
5
- model_path = './intent_classification/TinyBERT_106_V2' # can try other checkpoints
6
-
7
- tokenizer = BertTokenizer.from_pretrained(model_path)
8
- # model = BertForSequenceClassification.from_pretrained(model_path)
9
- model = AutoModelForSequenceClassification.from_pretrained(model_path, local_files_only=True)
10
- print(os.path.exists(model_path))
11
- print("TInyBERT model is ready to use")
12
-
13
-
14
- # for classification pipeline
15
- text_pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)
16
-
17
- # function to generate response
18
- def generate_response(user_query):
19
- response = text_pipeline(user_query)
20
-
21
- # example of response: [{'label': 'LABEL_4', 'score': 0.9997817873954773}]
22
- label_name = response[0].get('label')
23
- score = response[0].get('score')
24
-
25
- # label for each math topic based on label_name
26
- topic_label='NA'
27
-
28
- match label_name:
29
- case "LABEL_0":
30
- topic_label='RAG'
31
-
32
- case "LABEL_1":
33
- topic_label = 'Neo4j'
34
-
35
- return topic_label, score
36
 
37
 
38
  def get_dir():
39
  return os.getcwd()
40
 
41
- print(generate_response("Procedure to withdraw"))
42
  get_dir()
 
2
 
3
  from transformers import BertTokenizer, BertForSequenceClassification,TextClassificationPipeline, AutoModelForSequenceClassification
4
  # Load tokenizer and model from the fine-tuned directory
5
+ # model_path = './intent_classification/TinyBERT_106_V2' # can try other checkpoints
6
+ #
7
+ # tokenizer = BertTokenizer.from_pretrained(model_path)
8
+ # # model = BertForSequenceClassification.from_pretrained(model_path)
9
+ # model = AutoModelForSequenceClassification.from_pretrained(model_path, local_files_only=True)
10
+ # print(os.path.exists(model_path))
11
+ # print("TInyBERT model is ready to use")
12
+ #
13
+ #
14
+ # # for classification pipeline
15
+ # text_pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)
16
+ #
17
+ # # function to generate response
18
+ # def generate_response(user_query):
19
+ # response = text_pipeline(user_query)
20
+ #
21
+ # # example of response: [{'label': 'LABEL_4', 'score': 0.9997817873954773}]
22
+ # label_name = response[0].get('label')
23
+ # score = response[0].get('score')
24
+ #
25
+ # # label for each math topic based on label_name
26
+ # topic_label='NA'
27
+ #
28
+ # match label_name:
29
+ # case "LABEL_0":
30
+ # topic_label='RAG'
31
+ #
32
+ # case "LABEL_1":
33
+ # topic_label = 'Neo4j'
34
+ #
35
+ # return topic_label, score
36
 
37
 
38
  def get_dir():
39
  return os.getcwd()
40
 
41
+ # print(generate_response("Procedure to withdraw"))
42
  get_dir()