update

app.py

@@ -1,25 +1,87 @@
 import gradio as gr
-from llama_cpp import Llama
 import os 
+from pathlib import Path
+import argparse
 model_file = "Yi-6B.q4_k_m.gguf"
 if not os.path.isfile("Yi-6B.q4_k_m.gguf"):
     os.system("wget -c https://huggingface.co/SamPurkis/Yi-6B-GGUF/resolve/main/Yi-6B.q4_k_m.gguf")
 
-llm = Llama(model_path=model_file)
+DEFAULT_MODEL_PATH = model_file
+parser = argparse.ArgumentParser()
+parser.add_argument("-m", "--model", default=DEFAULT_MODEL_PATH, type=Path, help="model path")
+parser.add_argument("--mode", default="chat", type=str, choices=["chat", "generate"], help="inference mode")
+parser.add_argument("-l", "--max_length", default=2048, type=int, help="max total length including prompt and output")
+parser.add_argument("-c", "--max_context_length", default=512, type=int, help="max context length")
+parser.add_argument("--top_k", default=0, type=int, help="top-k sampling")
+parser.add_argument("--top_p", default=0.7, type=float, help="top-p sampling")
+parser.add_argument("--temp", default=0.95, type=float, help="temperature")
+parser.add_argument("--repeat_penalty", default=1.1, type=float, help="penalize repeat sequence of tokens")
+parser.add_argument("-t", "--threads", default=0, type=int, help="number of threads for inference")
+parser.add_argument("--plain", action="store_true", help="display in plain text without markdown support")
+args = parser.parse_args()
 
-def generate_text(input_text):
-    output = llm(f"Human: {input_text} A:", max_tokens=512, stop=["Assistant:", "\n"], echo=True)
-    return output['choices'][0]['text']
+from pyllamacpp.model import Model
+model = Model(model_path=model_file)
 
-input_text = gr.inputs.Textbox(lines= 10, label="Enter your input text")
-output_text = gr.outputs.Textbox(label="Output text")
 
-description = "llama.cpp implementation in python [https://github.com/abetlen/llama-cpp-python]"
 
-examples = [
-    ["What is the capital of France? ", "The capital of France is Paris."],
-    ["Who wrote the novel 'Pride and Prejudice'?", "The novel 'Pride and Prejudice' was written by Jane Austen."],
-    ["What is the square root of 64?", "The square root of 64 is 8."]
-]
+def predict(input, system_prompt, chatbot, max_length, ctx_length, top_p, temperature, history):
+    chatbot.append((input, ""))
+    response = ""
+    history.append(input)
 
-gr.Interface(fn=generate_text, inputs=input_text, outputs=output_text, title="Llama Language Model", description=description, examples=examples).launch()
+    generation_kwargs = dict(
+        max_length=max_length,
+        max_context_length=ctx_length,
+        do_sample=temperature > 0,
+        top_k=40,
+        top_p=top_p,
+        temperature=temperature,
+        repetition_penalty=1.1,
+        num_threads=0,
+        stream=True,
+    )
+    for response_piece in model.generate(input):
+        response += response_piece
+        chatbot[-1] = (chatbot[-1][0], response)
+
+        yield chatbot, history
+
+    history.append(response)
+    yield chatbot, history
+
+
+def reset_user_input():
+    return gr.update(value="")
+
+
+def reset_state():
+    return [], []
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""<h1 align="center">01-Yi 6B</h1>""")
+
+    chatbot = gr.Chatbot()
+    with gr.Row():
+        with gr.Column(scale=4):
+            system_prompt = gr.Textbox(show_label=False, placeholder="system prompt ...", lines=2)
+            user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=6)
+            submitBtn = gr.Button("Submit", variant="primary")
+        with gr.Column(scale=1):
+            max_length = gr.Slider(0, 32048, value=args.max_length, step=1.0, label="Maximum Length", interactive=True)
+            ctx_length = gr.Slider(0, 4096, value=512, step=1.0, label="Maximum Context Length", interactive=True)
+            top_p = gr.Slider(0, 1, value=args.top_p, step=0.01, label="Top P", interactive=True)
+            temperature = gr.Slider(0, 1, value=args.temp, step=0.01, label="Temperature", interactive=True)
+            emptyBtn = gr.Button("Clear History")
+
+    history = gr.State([])
+
+    submitBtn.click(
+        predict, [user_input, system_prompt, chatbot, max_length, ctx_length, top_p, temperature, history], [chatbot, history], show_progress=True
+    )
+    submitBtn.click(reset_user_input, [], [user_input])
+
+    emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
+
+demo.queue().launch(share=False, inbrowser=True)

requirements.txt

@@ -1,2 +1,2 @@
 gradio
-llama-cpp-python
+git+https://github.com/zhangheli/pyllamacpp.git@main