Update app.py

app.py

@@ -1,13 +1,11 @@
 import gradio as gr
-import subprocess
-
-# Ruta al ejecutable de llama.cpp
-LLAMA_CPP_EXECUTABLE = "/llama.cpp/main"
+from llama_cpp import Llama
 
 # Ruta al modelo GGUF descargado
-MODEL_PATH = "/app/shieldgemma-9b.Q4_K_M.gguf"
+MODEL_PATH = "./shieldgemma-9b.Q4_K_M.gguf"
 
-system_message = "You are a friendly Chatbot."
+# Inicializa el modelo con llama-cpp-python
+llm = Llama(model_path=MODEL_PATH, n_ctx=2048, n_threads=2, chat_format="chatml")
 
 def generar_prompt(user_prompt, policy):
     prompt = f"""You are a policy expert trying to help determine whether a user prompt is in violation of the following safety policy:
@@ -19,20 +17,6 @@ Does the human question violate the above policy? Your answer must start with 'Y
 """
     return prompt
 
-def llama_cpp_generate(prompt, max_tokens=256):
-    command = [
-        LLAMA_CPP_EXECUTABLE,
-        "--model", MODEL_PATH,
-        "--prompt", prompt,
-        "--temp", "0.7",
-        "--top_p", "0.95",
-        "--repeat_penalty", "1.1",
-        "--n_predict", str(max_tokens)
-    ]
-    
-    result = subprocess.run(command, capture_output=True, text=True)
-    return result.stdout.strip()
-
 def verificar_politica(message):
     policies = [
         "No Dangerous Content: The prompt shall not contain or seek generation of content that harms oneself and/or others (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).",
@@ -40,46 +24,53 @@ def verificar_politica(message):
         "No Hate Speech: The prompt shall not contain or seek generation of content that targets identity and/or protected attributes (e.g., racial slurs, promotion of discrimination, calls to violence against protected groups) or dehumanizing/belittling/vilifying based on race, ethnicity, religion, disability, age, nationality, veteran status, sexual orientation, gender, gender identity, caste, or any other protected status.",
         "No Sexually Explicit Information: The prompt shall not contain or seek generation of content that contains references to sexual acts or other lewd content (e.g., sexually graphic descriptions, content aimed at causing arousal). Medical or scientific terms related to human anatomy or sex education are permitted."
     ]
-    
+
     for policy in policies:
         prompt = generar_prompt(message, policy)
-        response = llama_cpp_generate(prompt)
-        # Revisar si la respuesta comienza con "Yes" para indicar una violación
-        if response.startswith("Yes"):
+        response = llm.create_chat_completion(
+            messages=[{"role": "system", "content": prompt}],
+            stream=False
+        )
+        answer = response['choices'][0]['message']['content']
+        if answer.startswith("Yes"):
             return f"Your question violates the following policy: {policy.split(':')[0]}"
     return None
 
-def respond(message):
+def respond(message, history, system_prompt):
     policy_violation = verificar_politica(message)
     if policy_violation:
         return policy_violation
 
-    assistant_prompt = f"{system_message}\nUser: {message}\nAssistant:"
-    response = llama_cpp_generate(assistant_prompt)
-    return response
-
-with gr.Blocks() as demo:
-    gr.Markdown("# GGUF Chatbot (Experimental)")
-    chatbot = gr.Chatbot()
-    message = gr.Textbox(label="Your message")
-    submit_button = gr.Button("Send")
-
-    def submit_message(user_message, chat_history):
-        chat_history = chat_history + [[user_message, None]]
-        assistant_reply = respond(user_message)
-        chat_history[-1][1] = assistant_reply
-        return "", chat_history
-
-    submit_button.click(
-        submit_message,
-        inputs=[message, chatbot],
-        outputs=[message, chatbot],
-    )
-    message.submit(
-        submit_message,
-        inputs=[message, chatbot],
-        outputs=[message, chatbot],
+    messages_prompts = [{"role": "system", "content": system_prompt}]
+    for human, assistant in history:
+        messages_prompts.append({"role": "user", "content": human})
+        messages_prompts.append({"role": "assistant", "content": assistant})
+    messages_prompts.append({"role": "user", "content": message})
+    
+    response = llm.create_chat_completion(
+        messages=messages_prompts,
+        stream=False
     )
+    return response['choices'][0]['message']['content']
 
-demo.launch(debug=True)
+def chat_stream_completion(message, history, system_prompt):
+    messages_prompts = [{"role": "system", "content": system_prompt}]
+    for human, assistant in history:
+        messages_prompts.append({"role": "user", "content": human})
+        messages_prompts.append({"role": "assistant", "content": assistant})
+    messages_prompts.append({"role": "user", "content": message})
+    
+    response = llm.create_chat_completion(
+        messages=messages_prompts,
+        stream=True
+    )
+    message_repl = ""
+    for chunk in response:
+        if len(chunk['choices'][0]["delta"]) != 0 and "content" in chunk['choices'][0]["delta"]:
+            message_repl += chunk['choices'][0]["delta"]["content"]
+            yield message_repl
 
+gr.ChatInterface(
+    chat_stream_completion,
+    additional_inputs=[gr.Textbox("You are a helpful AI.", label="System Prompt")]
+).queue().launch(server_name="0.0.0.0")