import os
import gradio as gr
import copy
import time
import llama_cpp
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

zefiro = Llama(
    model_path=hf_hub_download(
        repo_id="giux78/zefiro-7b-beta-ITA-v0.1-GGUF",
        filename="zefiro-7b-beta-ITA-v0.1-q4_0.gguf",
    ),
    n_ctx=4086,
) 


history = []

def generate_text(message, history):
    temp = ""
    input_prompt = "Chiedi a zefiro"
    for interaction in history:
        input_prompt += "[|Umano|] " + interaction[0] + "\n"
        input_prompt += "[|Assistente|]" + interaction[1]
    
    input_prompt += "[|Umano|] " + message + "\n[|Assistente|]"

    print(input_prompt)
    
    output = zefiro(input_prompt,
    temperature= 0.15,
    top_p= 0.1,
    top_k= 40, 
    repeat_penalty= 1.1,
    max_tokens= 1024,
    stop= [
        "[|Umano|]",
        "[|Assistente|]",
    ],
    stream= True)
    
    for out in output:
        stream = copy.deepcopy(out)
        temp += stream["choices"][0]["text"]
        yield temp

    history = ["init", input_prompt]



with gr.Blocks() as demo:
    with gr.Tab('zefiro'):
         gr.ChatInterface(
         generate_text,
         title="zefiro-7b-v01 running on CPU (quantized Q4_K)",
         description="This is a quantized version of zefiro-7b-v01 running on CPU (very slow). It is less powerful than the original version, but it can even run on the free tier of huggingface.",
         examples=[
             "Dammi 3 idee di ricette che posso fare con i pistacchi",
             "Prepara un piano di esercizi da poter fare a casa",
             "Scrivi una poesia su una giornato di pioggia"
         ],
         cache_examples=False,
         retry_btn=None,
         undo_btn="Delete Previous",
         clear_btn="Clear",
     )
        
demo.queue(concurrency_count=1, max_size=5)
demo.launch()