Update app.py

app.py

@@ -2,10 +2,10 @@ from threading import Thread
 
 import torch
 import gradio as gr
-from transformers import pipeline,AutoTokenizer, AutoModelForCausalLM, BertTokenizer, BertForSequenceClassification, StoppingCriteria, StoppingCriteriaList
+from transformers import pipeline,AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel, PeftConfig
 import re
-from kobert_transformers import get_tokenizer
+
 
 torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 print("Running on device:", torch_device)
@@ -33,29 +33,6 @@ model.eval()
 model.config.use_cache = True
 
 
-mbti_bert_model_name = "Lanvizu/fine-tuned-klue-bert-base_model_11"
-mbti_bert_model = BertForSequenceClassification.from_pretrained(mbti_bert_model_name)
-mbti_bert_model.eval()
-mbti_bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-
-bert_model_name = "ldhldh/bert_YN_small"
-bert_model = BertForSequenceClassification.from_pretrained(bert_model_name)
-bert_model.eval()
-bert_tokenizer = get_tokenizer()
-
-
-def mbti_classify(x):
-    classifier = pipeline("text-classification", model=mbti_bert_model, tokenizer=mbti_bert_tokenizer, return_all_scores=True)
-    result = classifier([x])
-    return result[0]
-
-
-def classify(x):
-    input_list =  bert_tokenizer.batch_encode_plus([x], truncation=True, padding=True, return_tensors='pt')
-    input_ids = input_list['input_ids'].to(bert_model.device)
-    attention_masks = input_list['attention_mask'].to(bert_model.device)
-    outputs =  bert_model(input_ids, attention_mask=attention_masks, return_dict=True)
-    return outputs.logits.argmax(dim=1).cpu().tolist()[0]
 
 def gen(x, top_p, top_k, temperature, max_new_tokens, repetition_penalty):
     gened = model.generate(
@@ -101,8 +78,7 @@ with gr.Blocks() as demo:
             )
             model_output = gr.Textbox(label="Model output", lines=10, interactive=False)
             button_submit = gr.Button(value="Submit")
-            button_bert = gr.Button(value="bert_Sumit")
-            button_mbti_bert = gr.Button(value="mbti_bert_Sumit")
+
         with gr.Column(scale=1):
             max_new_tokens = gr.Slider(
                 minimum=1, maximum=200, value=20, step=1, interactive=True, label="Max New Tokens",
@@ -121,6 +97,5 @@ with gr.Blocks() as demo:
             )
     
     button_submit.click(gen, [user_text, top_p, top_k, temperature, max_new_tokens, repetition_penalty], model_output)
-    button_bert.click(classify, [user_text], model_output)
-    button_mbti_bert.click(mbti_classify, [user_text], model_output)
+    
     demo.queue(max_size=32).launch(enable_queue=True)