Update app.py

app.py

@@ -6,8 +6,8 @@ from transformers import GenerationConfig
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 #Base Model 및 Lora Model 선택
-base_model = "beomi/KoAlpaca-Polyglot-5.8B"
-lora_weights = 'KimSHine/Scenario_Koalpaca_5.8B-lora'
+base_model = "EleutherAI/polyglot-ko-5.8b" # "beomi/KoAlpaca-Polyglot-5.8B"
+lora_weights = "KimSHine/Scenario_Koalpaca_v0_5.8B-lora" # 'KimSHine/Scenario_Koalpaca_5.8B-lora'
 load_8bit = True
 
 # Base Model Tokenizer
@@ -31,6 +31,7 @@ model1.config.pad_token_id = 0
 model1.config.eos_token_id = 2
 
 """### LoRA Model 불러오기
+
 Fine Tuning한 Model
 """
 
@@ -44,35 +45,73 @@ model1.config.pad_token_id = 0  # unk
 model1.config.bos_token_id = 0
 model1.config.eos_token_id = 2
 
-def yeollm_text(instruction, summary):
-
-    prompt = f"""아래는 작업을 설명하는 지시문과 대본을 생성하는데 참고할 줄거리와 짝을 이루는 예제입니다. 요청을 적절히 만족하는 대본을 작성하세요.
+val_dict = {"다큐멘터리": {
+                      'instruction' : "다큐멘터리 형식의 대본으로 만들어줘.",
+                      'temperature' :0.3,
+                      'top_p': 0.95,
+                      'top_k':40,
+                      'max_new_tokens':2048,
+                      'no_repeat_ngram_size': 5,
+                      'do_sample' : True,
+                      'num_beams' : 5},
+            "인터뷰": {
+                 'instruction' : "줄거리를 참고해서 인터뷰 형식의 대본을 만드시오. 인터뷰는 인터뷰어와 인터뷰이의 대화이며 인터뷰어가 질문을 하고 인터뷰이가 대답을 하는 형식입니다. 같은 말을 반복하지 마시오.",
+                      'temperature' :0.7,
+                      'top_p': 0.95,
+                      'top_k':40,
+                      'max_new_tokens':2048,
+                      'no_repeat_ngram_size': 5,
+                      'do_sample' : True,
+                      'num_beams' : 5},
+            "뉴스": {
+                 'instruction' : "뉴스  형식의 대본으로 만들어줘.",
+                      'temperature' :0.3,
+                      'top_p': 0.95,
+                      'top_k':40,
+                      'max_new_tokens':2048,
+                      'no_repeat_ngram_size': 5,
+                      'do_sample' : True,
+                      'num_beams' : 5},
+            "현대드라마": {
+                 'instruction' : "드라마 형식의 대본으로 만들어줘.",
+                      'temperature' :0.3,
+                      'top_p': 0.95,
+                      'top_k':40,
+                      'max_new_tokens':2048,
+                      'no_repeat_ngram_size': 5,
+                      'do_sample' : True,
+                      'num_beams' : 5},
+            "사극": {
+                 'instruction' : "사극 형식의 대본으로 만들어줘.",
+                      'temperature' :0.3,
+                      'top_p': 0.95,
+                      'top_k':40,
+                      'max_new_tokens':2048,
+                      'no_repeat_ngram_size': 5,
+                      'do_sample' : True,
+                      'num_beams' : 5}
+            }
+
+def yeollm_text(selected_value, summary):
+
+    prompt = f"""아래는 작업을 설명하는 지시문과 대본을 생성하는데 참고할 줄거리입니다.\n
     ### 지시문:
-    {instruction}
+    {val_dict[selected_value]['instruction']}
     ### 줄거리:
     {summary}
     ### 대본:
     """
 
-    temperature = 0.3
-    top_p = 0.95
-    top_k = 40
-    max_new_tokens = 512 #2048
-    no_repeat_ngram_size = 5    # 3개 이상의 토큰이 반복될 경우 확률을 0으로 만듦
-    ## greed search, beam search의 결과는 바뀌지 않음 (물론 최종 결과는 바뀜, 중간 sample 만드는 것은 동일하다는 것)
-    do_sample = True    ## True : random, False(default) : Greedy Search
-    num_beams = 5   ## do_sample 이 false일 때 여기에 값이 있으면, beam search
-
     inputs = tokenizer1(prompt, return_tensors="pt")
     input_ids = inputs["input_ids"].to(DEVICE)
 
     generation_config = GenerationConfig(
-        do_sample = do_sample,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
+        do_sample = val_dict[selected_value]['do_sample'],
+        temperature=val_dict[selected_value]['temperature'],
+        top_p=val_dict[selected_value]['top_p'],
+        top_k=val_dict[selected_value]['top_k'],
         pad_token_id = 0,   # pad token 추가
-        no_repeat_ngram_size = no_repeat_ngram_size,
+        no_repeat_ngram_size = val_dict[selected_value]['no_repeat_ngram_size'],
         # num_beams=num_beams,
         # **kwargs,
     )
@@ -84,7 +123,7 @@ def yeollm_text(instruction, summary):
             generation_config=generation_config,
             return_dict_in_generate=True,
             output_scores=True,
-            max_new_tokens=max_new_tokens,
+            max_new_tokens=val_dict[selected_value]['max_new_tokens'],
         )
     s = generation_output.sequences[0]
     output = tokenizer1.decode(s)
@@ -104,9 +143,9 @@ max_tokens = 2048
 temperature = 0.3
 Top_p = 1
 
-def davinci_text(instruction, summary):
+def davinci_text(selected_value, summary):
     prompt = f"""
-    아래의 줄거리를 보고 {instruction}
+    줄거리를 참고해서 {val_dict[selected_value]['instruction']} 형식의 대본을 만들어줘.
     ### 줄거리:
     {summary}
     ### 대본:
@@ -121,7 +160,6 @@ def davinci_text(instruction, summary):
     )
     return response.choices[0].text.strip()
 
-
 """## gpt 3.5 turbo 불러오기"""
 
 import openai
@@ -134,10 +172,10 @@ temperature = 0.3
 Top_p = 1
 
 
-def gpt_text(instruction, summary):
+def gpt_text(selected_value, summary):
   prompt =  f"""
     ### 지시문:
-    {instruction}
+    줄거리를 참고해서 {val_dict[selected_value]['instruction']} 형식의 대본을 만들어줘.
     ### 줄거리:
     {summary}
     ### 대본:
@@ -155,7 +193,7 @@ def gpt_text(instruction, summary):
   for choice in response["choices"]:
       content = choice["message"]["content"]
 
-  return content
+  return content.lstrip()
 
 """# gradio"""
 
@@ -164,8 +202,9 @@ import gradio as gr
 generator1 = gr.Interface(
     fn=yeollm_text,
     inputs=[
-        gr.inputs.Textbox(label="Instruction", placeholder="줄거리를 참고해서 현대 드라마 형식의 대본을 만들어줘"),
-        gr.inputs.Textbox(label="Summary",placeholder="대본으로 바꾸고 싶은 줄거리")
+        gr.Dropdown(["다큐멘터리", "인터뷰", "뉴스", '현대드라마', '사극'], label="형식"),
+        #gr.inputs.Textbox(label="Instruction",placeholder="줄거리를 참고해서 현대 드라마 형식의 대본을 만들어줘"),
+        gr.inputs.Textbox(label="Summary",placeholder="대본으로 바꾸고 싶은 줄거리"),
     ],
     outputs=gr.outputs.Textbox(label="Yeollm Scenario"),
     title="Yeollm Scenario Generation",
@@ -176,8 +215,8 @@ generator1 = gr.Interface(
 generator2 = gr.Interface(
     fn=davinci_text,
     inputs=[
-        gr.inputs.Textbox(label="Instruction", placeholder="줄거리를 참고해서 현대 드라마 형식의 대본을 만들어줘"),
-        gr.inputs.Textbox(label="Summary",placeholder="대본으로 바꾸고 싶은 줄거리")
+        gr.Dropdown(["다큐멘터리", "인터뷰", "뉴스", '현대드라마', '사극'], label="형식"),
+        gr.inputs.Textbox(label="Summary")
     ],
     outputs=gr.outputs.Textbox(label="Davinci Scenario"),
     title="Davinci Generation",
@@ -188,8 +227,8 @@ generator2 = gr.Interface(
 generator3 = gr.Interface(
     fn=gpt_text,
     inputs=[
-        gr.inputs.Textbox(label="Instruction", placeholder="줄거리를 참고해서 현대 드라마 형식의 대본을 만들어줘"),
-        gr.inputs.Textbox(label="Summary",placeholder="대본으로 바꾸고 싶은 줄거리")
+        gr.Dropdown(["다큐멘터리", "인터뷰", "뉴스", '현대드라마', '사극'], label="형식"),
+        gr.inputs.Textbox(label="Summary")
     ],
     outputs=gr.outputs.Textbox(label="GPT Scenario"),
     title="GPT Generation",
@@ -197,4 +236,4 @@ generator3 = gr.Interface(
     theme="huggingface"
 )
 
-gr.Parallel(generator1, generator2, generator3).launch()
+gr.Parallel(generator1, generator2, generator3).launch()