Update app.py

app.py

@@ -1,95 +1,62 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
-from llama_cpp_agent.providers import LlamaCppPythonProvider
-from llama_cpp_agent.chat_history import BasicChatHistory
-from llama_cpp_agent.chat_history.messages import UserMessage, AssistantMessage
-
-hf_hub_download(
-    repo_id="mradermacher/Arsh-llm-GGUF",
-    filename="Arsh-llm.Q4_K_M.gguf",
-    local_dir="./models"
-)
-
-llm = Llama(
-    model_path="./models/Arsh-llm.Q4_K_M.gguf",
-    n_batch=512,
-    n_ctx=8192,
-    verbose=False
-)
-
-provider = LlamaCppPythonProvider(llm)
-
-agent = LlamaCppAgent(
-    provider,
-    system_prompt="You are Arsh, a helpful assistant by Arshia Afshani. You should answer the user carefully.",
-    predefined_messages_formatter_type=MessagesFormatterType.CHATML,
-    debug_output=False
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+
+# Load model and tokenizer
+model_name = "Arsh-llm"  # نام مدل شما در Hugging Face
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+
+# Create pipeline
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1
 )
 
 def respond(message, chat_history, system_message, max_tokens, temperature, top_p, top_k, repeat_penalty):
-    settings = provider.get_provider_default_settings()
-    settings.temperature = temperature
-    settings.top_k = top_k
-    settings.top_p = top_p
-    settings.max_tokens = max_tokens
-    settings.repeat_penalty = repeat_penalty
-    settings.stream = True
-
-    if chat_history is None:
-        chat_history = []
-
-    messages = BasicChatHistory()
-
-    for msg in chat_history:
-        if msg["role"] == "user":
-            messages.add_message(UserMessage(msg["content"]))
-        elif msg["role"] == "assistant":
-            messages.add_message(AssistantMessage(msg["content"]))
-
-    stream = agent.get_chat_response(
-        message,
-        chat_history=messages,
-        llm_sampling_settings=settings,
-        returns_streaming_generator=True,
-        print_output=False
+    # Prepare prompt
+    prompt = f"{system_message}\n\nUser: {message}\nAssistant:"
+    
+    # Generate response
+    output = pipe(
+        prompt,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repeat_penalty,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     )
-
-    response = ""
-    for token in stream:
-        response += token
-        yield response
+    
+    response = output[0]['generated_text'].split("Assistant:")[-1].strip()
+    
+    # Update chat history
+    chat_history.append((message, response))
+    
+    return "", chat_history
 
 with gr.Blocks() as demo:
     gr.Markdown("# Arsh-LLM Demo")
 
     with gr.Row():
         with gr.Column():
-            system_msg = gr.Textbox("You are Arsh, a helpful assistant by Arshia Afshani. You should answer the user carefully.", label="System Message", interactive=True)
+            system_msg = gr.Textbox("You are Arsh, a helpful assistant by Arshia Afshani. You should answer the user carefully.", 
+                                  label="System Message", interactive=True)
             max_tokens = gr.Slider(1, 4096, value=2048, step=1, label="Max Tokens")
             temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
             top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
             top_k = gr.Slider(0, 100, value=40, step=1, label="Top-k")
             repeat_penalty = gr.Slider(0.0, 2.0, value=1.1, step=0.1, label="Repetition Penalty")
 
-    chatbot = gr.Chatbot(
-        bubble_full_width=False,
-        show_copy_button=True
-    )
+    chatbot = gr.Chatbot(bubble_full_width=False, show_copy_button=True)
+    msg = gr.Textbox(label="Your Message")
+    clear = gr.Button("Clear")
 
-    chat_interface = gr.ChatInterface(
-        respond,
-        additional_inputs=[
-            system_msg,
-            max_tokens,
-            temperature,
-            top_p,
-            top_k,
-            repeat_penalty
-        ],
-        chatbot=chatbot
-    )
+    msg.submit(respond, [msg, chatbot, system_msg, max_tokens, temperature, top_p, top_k, repeat_penalty], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
     demo.launch()