Update app.py

app.py

@@ -1,5 +1,5 @@
 from threading import Thread
-
+from llama_cpp import Llama
 import torch
 import gradio as gr
 import re
@@ -8,56 +8,18 @@ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 print("Running on device:", torch_device)
 print("CPU threads:", torch.get_num_threads())
 
-peft_model_id = "ldhldh/1.3_40kstep" 
-
-#peft_model_id = "ldhldh/polyglot-ko-1.3b_lora_big_tern_30kstep" 
-# 20k or 30k 
-#18k > 상대의 말까지 하는 이슈가 있음 
-#8k > 약간 아쉬운가? 
-
-base_model = AutoModelForCausalLM.from_pretrained("Skyranch/KoAlpaca-Polyglot-12.8B-ggml-model-f16", hf=True, model_type='gpt_neox')
-tokenizer = AutoTokenizer.from_pretrained("beomi/KoAlpaca-Polyglot-12.8B")
-
-#base_model = AutoModelForCausalLM.from_pretrained("EleutherAI/polyglot-ko-3.8b")
-#tokenizer = AutoTokenizer.from_pretrained("EleutherAI/polyglot-ko-3.8b")
-base_model.eval()
-base_model.config.use_cache = True
-
-
-#model = PeftModel.from_pretrained(base_model, peft_model_id, device_map="auto")
-#model.eval()
-#model.config.use_cache = True
-
-def gen(x, top_p, top_k, temperature, max_new_tokens, repetition_penalty):
-    gened = base_model.generate(
-        **tokenizer(
-            f"{x}",
-            return_tensors='pt',
-            return_token_type_ids=False
-        ),
-        #bad_words_ids = bad_words_ids ,
-        max_new_tokens=max_new_tokens,
-        min_new_tokens = 5,
-        exponential_decay_length_penalty = (max_new_tokens/2, 1.1),
-        top_p=top_p,
-        top_k=top_k,
-        temperature = temperature,
-        early_stopping=True,
-        do_sample=True,
-        eos_token_id=2,
-        pad_token_id=2,
-        #stopping_criteria = stopping_criteria,
-        repetition_penalty=repetition_penalty,
-        no_repeat_ngram_size = 2
-    )
+llm = Llama(model_path = 'Llama-2-ko-7B-chat-gguf-q4_0.bin',
+            n_ctx=1024,
+      )
 
-    model_output = tokenizer.decode(gened[0])
-    return model_output
+def gen(x, max_new_tokens):
+    output = llm(f"Q: {x} A: ", max_tokens=1024, stop=["Q:", "\n"], echo=True)
+    
+    return output['choices'][0]['text'].replace('▁',' ')
 
 def reset_textbox():
     return gr.update(value='')
 
-
 with gr.Blocks() as demo:
     duplicate_link = "https://huggingface.co/spaces/beomi/KoRWKV-1.5B?duplicate=true"
     gr.Markdown(
@@ -67,7 +29,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=4):
             user_text = gr.Textbox(
-                placeholder='\\nfriend: 우리 여행 갈래? \\nyou:',
+                placeholder='우리 여행 갈래?',
                 label="User input"
             )
             model_output = gr.Textbox(label="Model output", lines=10, interactive=False)
@@ -76,18 +38,6 @@ with gr.Blocks() as demo:
             max_new_tokens = gr.Slider(
                 minimum=1, maximum=200, value=20, step=1, interactive=True, label="Max New Tokens",
             )
-            top_p = gr.Slider(
-                minimum=0.05, maximum=1.0, value=0.8, step=0.05, interactive=True, label="Top-p (nucleus sampling)",
-            )
-            top_k = gr.Slider(
-                minimum=5, maximum=100, value=30, step=5, interactive=True, label="Top-k (nucleus sampling)",
-            )
-            temperature = gr.Slider(
-                minimum=0.1, maximum=2.0, value=0.5, step=0.1, interactive=True, label="Temperature",
-            )
-            repetition_penalty = gr.Slider(
-                minimum=1.0, maximum=3.0, value=1.2, step=0.1, interactive=True, label="repetition_penalty",
-            )
     
-    button_submit.click(gen, [user_text, top_p, top_k, temperature, max_new_tokens, repetition_penalty], model_output)
+    button_submit.click(gen, [user_text, max_new_tokens], model_output)
     demo.queue(max_size=32).launch(enable_queue=True)