Update app.py

app.py

@@ -1,11 +1,5 @@
 import gradio as gr
-import selfies as sf
 from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent
-from llama_cpp_agent.providers import LlamaCppPythonProvider
-from llama_cpp_agent.chat_history import BasicChatHistory
-from llama_cpp_agent.chat_history.messages import Roles
-from llama_cpp_agent import MessagesFormatterType
 
 css = """
 .message-row {
@@ -26,34 +20,37 @@ css = """
 """
 
 def respond(encoded_smiles, max_tokens, temperature, top_p, top_k):
-    model_name = "model.gguf"
-    llm = Llama(model_name)
-    provider = LlamaCppPythonProvider(llm)
-    chat_template = MessagesFormatterType.CHATML
-    
-    agent = LlamaCppAgent(
-        provider,
-        predefined_messages_formatter_type=chat_template,
-        debug_output=True
-    )
+    try:
+        # Load the Llama model
+        model_name = "model.gguf"
+        llm = Llama(model_name)  # Khởi tạo đối tượng Llama với tệp mô hình
+        
+        # Set generation settings
+        settings = {
+            "max_new_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "do_sample": True,
+        }
 
-    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.stream = False
+        # Tokenize the input
+        input_ids = llm.tokenizer(encoded_smiles, return_tensors='pt').input_ids
 
-    prompt = f"{encoded_smiles}"
-    input_ids = agent.tokenizer(prompt, return_tensors='pt', truncation=False).input_ids.cuda()
-    outputs = agent.llm.generate(input_ids=input_ids)
-    output1 = agent.tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]
+        # Generate the output
+        outputs = llm.generate(input_ids=input_ids, **settings)
 
-    first_inst_index = output1.find("[/INST]")
-    second_inst_index = output1.find("[/IN", first_inst_index + len("[/INST]") + 1)
-    predicted_selfies = output1[first_inst_index + len("[/INST]") : second_inst_index].strip()
+        # Decode the output tokens to text
+        output_text = llm.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    return {'input': encoded_smiles, 'predict': predicted_selfies}
+        # Extract the predicted selfies from the output text
+        first_inst_index = output_text.find("[/INST]")
+        second_inst_index = output_text.find("[/IN", first_inst_index + len("[/INST]") + 1)
+        predicted_selfies = output_text[first_inst_index + len("[/INST]"): second_inst_index].strip()
+
+        return {'input': encoded_smiles, 'predict': predicted_selfies}
+    except Exception as e:
+        return {'error': str(e)}
 
 demo = gr.Interface(
     fn=respond,
@@ -61,20 +58,8 @@ demo = gr.Interface(
         gr.Textbox(label="Encoded SMILES"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=1.0, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=1.0,
-            step=0.05,
-            label="Top-p",
-        ),
-        gr.Slider(
-            minimum=0,
-            maximum=100,
-            value=50,
-            step=1,
-            label="Top-k",
-        )
+        gr.Slider(minimum=0.1, maximum=1.0, value=1.0, step=0.05, label="Top-p"),
+        gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-k")
     ],
     outputs=gr.JSON(label="Results"),
     theme=gr.themes.Soft(primary_hue="violet", secondary_hue="violet", neutral_hue="gray", font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]).set(