Update helper.py

helper.py

@@ -1,200 +1,96 @@
 import os
-import subprocess
-from huggingface_hub import hf_hub_download, list_repo_files
 import gradio as gr
 from typing import Callable
 import base64
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from threading import Thread
-from transformers import TextIteratorStreamer
+from openai import OpenAI
 
+def get_fn(model_name: str, **model_kwargs):
+    """Create a chat function that uses the OpenAI-compatible endpoint."""   
 
-def get_fn(model_path: str, **model_kwargs):
-    """Create a chat function with the specified model."""
-    
-    # Initialize tokenizer and model
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    
-    # Simple flash-attention installation attempt
-    try:
-        subprocess.run(
-            'pip install flash-attn --no-build-isolation',
-            env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
-            shell=True,
-            check=True
-        )
-        # Try loading model with flash attention
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            device_map="auto",
-            quantization_config=quantization_config,
-            attn_implementation="flash_attention_2",
-        )
-    except Exception as e:
-        print(f"Flash Attention failed, falling back to default attention: {str(e)}")
-        # Fallback to default attention implementation
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            torch_dtype= torch.bfloat16,
-            device_map="auto",
-            quantization_config=quantization_config,
-        )
+    OPENAI_API_KEY = "-"
+    client = OpenAI(
+    base_url=" http://192.222.58.60:8000/v1", 
+    api_key="tela",
+    )
 
     def predict(
         message: str,
         history,
         system_prompt: str,
         temperature: float,
-        max_new_tokens: int,
-        top_k: int,
-        repetition_penalty: float,
-        top_p: float
+        top_p: float,
+        max_tokens: int,
+        
     ):
         try:
-            # Format conversation with ChatML format
-            instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
+            messages = []
+            if system_prompt:
+                messages.append({"role": "system", "content": system_prompt})
             for user_msg, assistant_msg in history:
-                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'
-
-            streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-            enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
-            input_ids, attention_mask = enc.input_ids, enc.attention_mask
+                messages.append({"role": "user", "content": user_msg})
+                messages.append({"role": "assistant", "content": assistant_msg})
+            messages.append({"role": "user", "content": message})
 
-            # Truncate if needed
-            if input_ids.shape[1] > 8192:  # Using n_ctx from original
-                input_ids = input_ids[:, -8192:]
-                attention_mask = attention_mask[:, -8192:]
-
-            generate_kwargs = dict(
-                input_ids=input_ids.to(device),
-                attention_mask=attention_mask.to(device),
-                streamer=streamer,
-                do_sample=True,
+            response = openai.chat.completions.create(
+                model=model_name,
+                messages=messages,
                 temperature=temperature,
-                max_new_tokens=max_new_tokens,
-                top_k=top_k,
-                repetition_penalty=repetition_penalty,
-                top_p=top_p
+                top_p=top_p,
+                max_tokens=max_new_tokens,
+                n=1,
+                stream=True, 
+                response_format={"type": "text"},
             )
-
-            t = Thread(target=model.generate, kwargs=generate_kwargs)
-            t.start()
-
-            response_text = ""
-            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."
-
+            assistant_message = response.choices[0].message.content
+            yield assistant_message.strip()
         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 type(message) is str:
+    if isinstance(message, str):
         return message
-    elif type(message) is dict:
-        if message["files"] is not None and len(message["files"]) > 0:
+    elif isinstance(message, dict):
+        if message.get("files"):
             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["text"]},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": encoded_str,
-                        }
-                    },
-                ]
+                {"type": "text", "text": message.get("text", "")},
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": encoded_str,
+                    }
+                },
+            ]
         else:
-            content = message["text"]
+            content = message.get("text", "")
         return content
     else:
         raise NotImplementedError
 
-
-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 get_model_path(name: str = None, model_path: str = None) -> str:
-    """Get the local path to the model."""
+    """Get the model name to use with the endpoint."""
     if model_path:
         return model_path
-    
     if name:
-        if "/" in name:
-            return name  # Return HF model ID directly
-        else:
-            # You could maintain a mapping of friendly names to HF model IDs
-            model_mapping = {
-                # Add any default model mappings here
-                "example-model": "organization/model-name"
-            }
-            if name not in model_mapping:
-                raise ValueError(f"Unknown model name: {name}")
-            return model_mapping[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 Interface with similar styling and parameters."""
-    
-    model_path = get_model_path(name, model_path)
-    fn = get_fn(model_path, **kwargs)
+    """Create a Gradio ChatInterface."""
+    model_name = get_model_path(name, model_path)
+    fn = get_fn(model_name, **kwargs)
 
     interface = gr.ChatInterface(
         fn=fn,
@@ -206,10 +102,8 @@ def registry(name: str = None, model_path: str = None, **kwargs):
             ),
             gr.Slider(0, 1, 0.7, label="Temperature"),
             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