import gradio as gr
from huggingface_hub import InferenceClient
import torch
from transformers import LlamaTokenizer, LlamaForCausalLM

prompt_input = (
    "Below is an instruction that describes a task. "
    "Write a response that appropriately completes the request.\n\n"
    "### Instruction:\n\n{instruction}\n\n### Response:\n\n"
)
load_type = torch.float16
device = torch.device(0)


def generate_prompt(instruction, input=None):
    if input:
        instruction = instruction + '\n' + input
    return prompt_input.format_map({'instruction': instruction})


model_path = "ViraIntelligentDataMining/PersianLLaMA-13B-Instruct"
tokenizer = LlamaTokenizer.from_pretrained(model_path)
base_model = LlamaForCausalLM.from_pretrained(
    model_path,
    torch_dtype=load_type,
    device_map='cpu',
).cuda()

model_vocab_size = base_model.get_input_embeddings().weight.size(0)
tokenizer_vocab_size = len(tokenizer)
if model_vocab_size != tokenizer_vocab_size:
    base_model.resize_token_embeddings(tokenizer_vocab_size)


def generate_answer(base_model, instruction, input=None):
    generation_config = dict(
        temperature=0.5,
        top_k=40,
        top_p=0.9,
        repetition_penalty=1.1,
        max_new_tokens=1024)
    input_text = generate_prompt(instruction, input)
    inputs = tokenizer(input_text, return_tensors="pt")
    generation_output = base_model.generate(
        input_ids=inputs["input_ids"].to(device),
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        **generation_config)
    s = generation_output[0]
    output = tokenizer.decode(s, skip_special_tokens=True)
    response = output.split("### Response:")[1].strip()
    return response


instruction = "تصور کنید در حال نوشتن داستانی درباره یک شهر که تمام ساکنانش ربات هستند. این ربات‌ها توانایی‌های ویژه‌ای دارند که زندگی روزمره آن‌ها را از انسان‌ها متمایز می‌کند. شرح دهید که این ربات‌ها چگونه به کارهای روزانه خود می‌پردازند و چه چالش‌هایی با آن‌ها روبه‌رو هستند. همچنین، توضیح دهید که چگونه معماری شهر برای نیازهای خاص آن‌ها طراحی شده است."

response = generate_answer(base_model,
                           instruction=instruction,
                           input="")
print(response)
"""
در این شهر، همه ساکنان ربات هستند که دارای توانایی های منحصر به فرد هستند که زندگی روزمره آنها را از انسان ها متمایز می کند. هر روز صبح، ربات ها بیدار می شوند و برنامه های خود را برای روز تنظیم می کنند. برخی از آنها ممکن است برای کار بروند، در حالی که دیگران ممکن است برای انجام وظایف خانگی یا مراقبت از خانواده خود وقت صرف کنند. ربات ها بسیار کارآمد هستند و می توانند چندین کار را همزمان انجام دهند، بنابراین زمان زیادی برای استراحت ندارند. آنها همچنین به طور منظم برای نگهداری و تعمیر نیاز دارند، زیرا آنها مانند انسان ها مستعد خرابی هستند. بسیاری از ربات ها به دنبال سرگرمی هستند و ممکن است برای شرکت در فعالیت هایی مانند ورزش، خواندن یا نقاشی وقت صرف کنند. برخی از ربات ها حتی ممکن است برای یادگیری یا بهبود مهارت های خود در زمینه های مختلف مانند هنر، موسیقی یا آشپزی تلاش کنند. ربات ها همچنین به طور مداوم به پیشرفت های تکنولوژیکی جدید علاقه مند هستند و اغلب به دنبال راه حل های جدید برای مشکلات موجود در جامعه هستند. در این شهر، همه ساکنان ربات هستند که دارای توانایی های منحصر به فرد هستند که زندگی روزمره آنها را از انسان ها متمایز می کند.
"""


if __name__ == "__main__":
    demo.launch()