src/hf.py

import os.path

from transformers import BertTokenizer, BertForSequenceClassification,TextClassificationPipeline, AutoModelForSequenceClassification
# Load tokenizer and model from the fine-tuned directory
# model_path = './intent_classification/TinyBERT_106_V2'  # can try other checkpoints
#
# tokenizer = BertTokenizer.from_pretrained(model_path)
# # model = BertForSequenceClassification.from_pretrained(model_path)
# model = AutoModelForSequenceClassification.from_pretrained(model_path, local_files_only=True)
# print(os.path.exists(model_path))
# print("TInyBERT model is ready to use")
#
#
# # for classification pipeline
# text_pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)
#
# # function to generate response
# def generate_response(user_query):
#     response = text_pipeline(user_query)
#
#     # example of response: [{'label': 'LABEL_4', 'score': 0.9997817873954773}]
#     label_name = response[0].get('label')
#     score = response[0].get('score')
#
#     # label for each math topic based on label_name
#     topic_label='NA'
#
#     match label_name:
#         case "LABEL_0":
#             topic_label='RAG'
#
#         case "LABEL_1":
#             topic_label = 'Neo4j'
#
#     return topic_label, score


def get_dir():
    return os.getcwd()

# print(generate_response("Procedure to withdraw"))
get_dir()