Create main.py

main.py

@@ -0,0 +1,215 @@
+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
+
+
+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,
+        )
+
+    def predict(
+        message: str,
+        history,
+        system_prompt: str,
+        temperature: float,
+        max_new_tokens: int,
+        top_k: int,
+        repetition_penalty: float,
+        top_p: float
+    ):
+        try:
+            # Format conversation with ChatML format
+            instruction = '<|im_start|>system\n' + system_prompt + '\n<|im_end|>\n'
+            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
+
+            # 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,
+                temperature=temperature,
+                max_new_tokens=max_new_tokens,
+                top_k=top_k,
+                repetition_penalty=repetition_penalty,
+                top_p=top_p
+            )
+
+            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."
+
+        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:
+        return message
+    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["text"]},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": encoded_str,
+                        }
+                    },
+                ]
+        else:
+            content = message["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."""
+    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]
+    
+    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)
+
+    interface = gr.ChatInterface(
+        fn=fn,
+        additional_inputs_accordion=gr.Accordion("⚙️ Parameters", open=False),
+        additional_inputs=[
+            gr.Textbox(
+                "You are a helpful AI assistant.",
+                label="System prompt"
+            ),
+            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