import gradio as gr
from llama_cpp import Llama
import random

# Initialize model
llm = Llama(model_path="/stablelm-2-zephyr-1_6b-Q4_0.gguf", n_gpu_layers=0, seed=random.randint(1, 2**31))

def generate_response(user_message):
    encodeduserm = b"### Human: " + user_message.encode('utf-8') + b"\n### Assistant:"
    tokens = llm.tokenize(encodeduserm)
    output = b""
    count = 0

    for token in llm.generate(tokens, top_k=40, top_p=0.95, temp=0.72, repeat_penalty=1.1):
        text = llm.detokenize([token])
        output += text
        count += 1
        if count >= 500 or (token == llm.token_eos()):
            break
    return output.decode()

iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Type your message here..."),
    outputs="text",
    title="LLaMA Chat Interface",
    description="Enter your message and get a response from the LLaMA model."
)

iface.launch(share=True)