Update helper.py

helper.py

@@ -4,31 +4,39 @@ from typing import Callable
 import base64
 from openai import OpenAI
 
-def get_fn(model_name: str, **model_kwargs):
-    """Create a chat function that uses the OpenAI-compatible endpoint."""   
 
-    OPENAI_API_KEY = "-"
-    client = OpenAI(
-    base_url=" http://192.222.58.60:8000/v1", 
-    api_key="tela",
-    )
+
+def get_fn(model_path: str, **model_kwargs):
+    """Create a chat function with the specified model."""
+    
+    # instatiate a OpenAI client for a custom endpoint
+    try:
+        OPENAI_API_KEY = "-"
+        client = OpenAI(
+        base_url=" http://192.222.58.60:8000/v1", 
+        api_key="tela",
+        )
+
+    except Exception as e:
+        print(f"The api or base url were not definied: {str(e)}")
+
 
     def predict(
         message: str,
         history,
         system_prompt: str,
         temperature: float,
-        max_tokens: int,
-        top_p: float,  
+        max_new_tokens: int,
+        top_k: int,
+        repetition_penalty: float,
+        top_p: float
     ):
         try:
-            messages = []
-            if system_prompt:
-                messages.append({"role": "system", "content": system_prompt})
+            # Format conversation with ChatML format
+            instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
             for user_msg, assistant_msg in history:
-                messages.append({"role": "user", "content": user_msg})
-                messages.append({"role": "assistant", "content": assistant_msg})
-            messages.append({"role": "user", "content": message})
+                instruction += f'<|im_start|>user\n{user_msg}\n<|im_end|>\n<|im_start|>assistant\n{assistant_msg}\n<|im_end|>\n'
+            instruction += f'<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n'
 
             response = client.chat.completions.create(
                 model=model_name,
@@ -43,58 +51,95 @@ def get_fn(model_name: str, **model_kwargs):
 
             response_text = ""
             for chunk in response:
-                chunk_message = chunk.choices[0].delta.content
-                if chunk_message:
-                    response_text += chunk_message
+                streamer = chunk.choices[0].delta.content
+                for new_token in streamer:
+                    if new_token in ["<|endoftext|>", "<|im_end|>"]:
+                        break
+                    response_text += new_token
                     yield response_text.strip()
+
+            if not response_text.strip():
+                yield "I apologize, but I was unable to generate a response. Please try again."
+
         except Exception as e:
             print(f"Error during generation: {str(e)}")
             yield f"An error occurred: {str(e)}"
 
     return predict
 
+
 def get_image_base64(url: str, ext: str):
     with open(url, "rb") as image_file:
         encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
     return "data:image/" + ext + ";base64," + encoded_string
 
+
 def handle_user_msg(message: str):
-    if isinstance(message, str):
+    if type(message) is str:
         return message
-    elif isinstance(message, dict):
-        if message.get("files"):
+    elif type(message) is dict:
+        if message["files"] is not None and len(message["files"]) > 0:
             ext = os.path.splitext(message["files"][-1])[1].strip(".")
             if ext.lower() in ["png", "jpg", "jpeg", "gif", "pdf"]:
                 encoded_str = get_image_base64(message["files"][-1], ext)
             else:
                 raise NotImplementedError(f"Not supported file type {ext}")
             content = [
-                {"type": "text", "text": message.get("text", "")},
-                {
-                    "type": "image_url",
-                    "image_url": {
-                        "url": encoded_str,
-                    }
-                },
-            ]
+                    {"type": "text", "text": message["text"]},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": encoded_str,
+                        }
+                    },
+                ]
         else:
-            content = message.get("text", "")
+            content = message["text"]
         return content
     else:
         raise NotImplementedError
 
-def get_model_path(name: str = None, model_path: str = None) -> str:
-    """Get the model name to use with the endpoint."""
-    if model_path:
-        return model_path
-    if name:
-        return name
-    raise ValueError("Either name or model_path must be provided")
 
-def registry(name: str = None, model_path: str = None, **kwargs):
-    """Create a Gradio ChatInterface."""
-    model_name = get_model_path(name, model_path)
-    fn = get_fn(model_name, **kwargs)
+def get_interface_args(pipeline):
+    if pipeline == "chat":
+        inputs = None
+        outputs = None
+
+        def preprocess(message, history):
+            messages = []
+            files = None
+            for user_msg, assistant_msg in history:
+                if assistant_msg is not None:
+                    messages.append({"role": "user", "content": handle_user_msg(user_msg)})
+                    messages.append({"role": "assistant", "content": assistant_msg})
+                else:
+                    files = user_msg
+            if type(message) is str and files is not None:
+                message = {"text":message, "files":files}
+            elif type(message) is dict and files is not None:
+                if message["files"] is None or len(message["files"]) == 0:
+                    message["files"] = files
+            messages.append({"role": "user", "content": handle_user_msg(message)})
+            return {"messages": messages}
+
+        postprocess = lambda x: x
+    else:
+        # Add other pipeline types when they will be needed
+        raise ValueError(f"Unsupported pipeline type: {pipeline}")
+    return inputs, outputs, preprocess, postprocess
+
+
+def get_pipeline(model_name):
+    # Determine the pipeline type based on the model name
+    # For simplicity, assuming all models are chat models at the moment
+    return "chat"
+
+
+
+def registry(name: str = None, **kwargs):
+    """Create a Gradio Interface with similar styling and parameters."""
+    
+    fn = get_fn(name, **kwargs)
 
     interface = gr.ChatInterface(
         fn=fn,
@@ -105,9 +150,11 @@ def registry(name: str = None, model_path: str = None, **kwargs):
                 label="System prompt"
             ),
             gr.Slider(0, 1, 0.7, label="Temperature"),
-            gr.Slider(128, 4096, 1024, label="Max tokens"),
+            gr.Slider(128, 4096, 1024, label="Max new tokens"),
+            gr.Slider(1, 80, 40, label="Top K sampling"),
+            gr.Slider(0, 2, 1.1, label="Repetition penalty"),
             gr.Slider(0, 1, 0.95, label="Top P sampling"),
         ],
     )
-
-    return interface
+    
+    return interface