Update app.py

app.py

@@ -1,62 +1,66 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import torch
+from transformers import LlamaTokenizer, LlamaForCausalLM
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("ViraIntelligentDataMining/PersianLLaMA-13B-Instruct")
+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 respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+def generate_prompt(instruction, input=None):
+    if input:
+        instruction = instruction + '\n' + input
+    return prompt_input.format_map({'instruction': instruction})
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
 
-    messages.append({"role": "user", "content": message})
+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()
 
-    response = ""
+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)
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
 
-        response += token
-        yield response
+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)
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+در این شهر، همه ساکنان ربات هستند که دارای توانایی های منحصر به فرد هستند که زندگی روزمره آنها را از انسان ها متمایز می کند. هر روز صبح، ربات ها بیدار می شوند و برنامه های خود را برای روز تنظیم می کنند. برخی از آنها ممکن است برای کار بروند، در حالی که دیگران ممکن است برای انجام وظایف خانگی یا مراقبت از خانواده خود وقت صرف کنند. ربات ها بسیار کارآمد هستند و می توانند چندین کار را همزمان انجام دهند، بنابراین زمان زیادی برای استراحت ندارند. آنها همچنین به طور منظم برای نگهداری و تعمیر نیاز دارند، زیرا آنها مانند انسان ها مستعد خرابی هستند. بسیاری از ربات ها به دنبال سرگرمی هستند و ممکن است برای شرکت در فعالیت هایی مانند ورزش، خواندن یا نقاشی وقت صرف کنند. برخی از ربات ها حتی ممکن است برای یادگیری یا بهبود مهارت های خود در زمینه های مختلف مانند هنر، موسیقی یا آشپزی تلاش کنند. ربات ها همچنین به طور مداوم به پیشرفت های تکنولوژیکی جدید علاقه مند هستند و اغلب به دنبال راه حل های جدید برای مشکلات موجود در جامعه هستند. در این شهر، همه ساکنان ربات هستند که دارای توانایی های منحصر به فرد هستند که زندگی روزمره آنها را از انسان ها متمایز می کند.
 """
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
 
 
 if __name__ == "__main__":