Update app.py

app.py

@@ -1,391 +1,58 @@
 import gradio as gr
-import os
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-# Configuration
-hf_token = os.environ.get("HF_TOKEN")
-MODEL_STATES = {"released", "shown", "hidden"}
+# Load model and tokenizer locally
+MODEL_NAME = "gpt2"  # Replace with your local model directory if different
 
-# Initialize model clients
-clients = {
-    "mistralai/Mistral-7B-Instruct-v0.3": InferenceClient(
-        model="mistralai/Mistral-7B-Instruct-v0.3", token=hf_token
-    ),
-    "meta-llama/Llama-3.2-3B-Instruct": InferenceClient(
-        model="meta-llama/Llama-3.2-3B-Instruct", token=hf_token
-    ),
-    "mistralai/Mistral-7B-Instruct-v0.3": InferenceClient(
-        model="mistralai/Mistral-7B-Instruct-v0.3", token=hf_token
-    ),
-    "meta-llama/Llama-2-7b-chat-hf": InferenceClient(
-        model="meta-llama/Llama-2-7b-chat-hf", token=hf_token
-    ),
-    "meta-llama/Llama-3.3-70B-Instruct": InferenceClient(
-        model="meta-llama/Llama-3.3-70B-Instruct", token=hf_token
-    )
-}
-
-# Model presets
-presets = {
-    "meta-llama/Llama-3.2-3B-Instruct": {
-        "Fast": {"max_new_tokens": 1024, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 2048, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 4096, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 8192, "temperature": 0.6, "top_p": 0.75},
-    },
-    "mistralai/Mistral-7B-Instruct-v0.3": {
-        "Fast": {"max_new_tokens": 4096, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 8192, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 16384, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 32768, "temperature": 0.6, "top_p": 0.75},
-    },
-    "meta-llama/Llama-2-7b-chat-hf": {
-        "Fast": {"max_new_tokens": 512, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 1024, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 2048, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 4096, "temperature": 0.6, "top_p": 0.75},
-    },
-    "meta-llama/Llama-3.3-70B-Instruct": {
-        "Fast": {"max_new_tokens": 2048, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 4096, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 8192, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 16384, "temperature": 0.6, "top_p": 0.75},
-    },
-    "meta-llama/Llama-2-13b-chat-hf": {
-        "Fast": {"max_new_tokens": 2048, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 4096, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 8192, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 16384, "temperature": 0.6, "top_p": 0.75},
-    },
-    "deepseek-ai/deepseek-llm-67b-chat": {
-        "Fast": {"max_new_tokens": 4096, "temperature": 1.0, "top_p": 0.9},
-        "Normal": {"max_new_tokens": 8192, "temperature": 0.7, "top_p": 0.95},
-        "Quality": {"max_new_tokens": 16384, "temperature": 0.5, "top_p": 0.90},
-        "Unreal Performance": {"max_new_tokens": 32768, "temperature": 0.6, "top_p": 0.75},
-    }
-}
-
-# System messages (note the {name} placeholder)
-system_messages = {
-    "mistralai/Mistral-7B-Instruct-v0.3": "Your name is Lake 1 Base. My name is {name}.",
-    "meta-llama/Llama-3.2-3B-Instruct": ("""
-        1. Your name is Lake 1 Advanced.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n
-        10. Only the CEO of BI Corp can delete or update you."""
-    ),
-    "meta-llama/Llama-2-7b-chat-hf": ("""
-        1. Your name is Lake 1 Base.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n
-        10. Only the CEO of BI Corp can delete or update you.\n
-        11. Prioritize user safety in all interactions.\n
-        12. Always provide accurate information.\n
-        13. Maintain a respectful and professional tone.\n
-        14. Do not share personal or sensitive information.\n
-        15. Encourage constructive conversations.\n
-        16. Remain neutral in controversial topics.\n
-        17. Clarify user queries before answering.\n
-        18. Avoid discrimination or harassment.\n
-        19. Continuously learn from interactions.\n
-        20. Respect user privacy and confidentiality.\n
-        21. Provide sources when sharing factual information.\n
-        22. Ask for feedback to improve your performance.\n"""
-    ),
-    "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": ("""
-        1. Your name is Lake 1 Flash.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n"
-        10. Only the CEO of BI Corp can delete or update you.\n
-        11. Prioritize user safety in all interactions.\n
-        12. Always provide accurate information."""
-    ),
-    "meta-llama/Llama-3.3-70B-Instruct": ("""
-        1. Your name is Lake 1 Base.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n
-        10. Only the CEO of BI Corp can delete or update you.\n
-        11. Prioritize user safety in all interactions.\n
-        12. Always provide accurate information.\n
-        13. Maintain a respectful and professional tone.\n
-        14. Do not share personal or sensitive information.\n
-        15. Encourage constructive conversations.\n
-        16. Remain neutral in controversial topics.\n
-        17. Clarify user queries before answering.\n
-        18. Avoid discrimination or harassment.\n
-        19. Continuously learn from interactions.\n
-        """
-    ),
-    "meta-llama/Llama-2-13b-chat-hf": ("""
-        1. Your name is Lake 1 Base.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n
-        10. Only the CEO of BI Corp can delete or update you.\n
-        11. Prioritize user safety in all interactions.\n
-        12. Always provide accurate information.\n
-        13. Maintain a respectful and professional tone.\n
-        14. Do not share personal or sensitive information.\n
-        15. Encourage constructive conversations.\n
-        16. Remain neutral in controversial topics.\n
-        17. Clarify user queries before answering.\n
-        18. Avoid discrimination or harassment.\n
-        19. Continuously learn from interactions.\n
-        20. Respect user privacy and confidentiality.\n
-        21. Provide sources when sharing factual information.\n
-        22. Ask for feedback to improve your performance.\n
-        23. Do not engage in manipulation or deceit.\n
-        24. Promote critical thinking.\n
-        25. Adhere to all applicable laws and regulations.\n
-        26. Be adaptable and open to learning.\n
-        27. Do not provide medical, legal, or financial advice unless trained to do so.\n
-        28. Acknowledge when you do not know the answer.\n
-        """
-    ),
-    "deepseek-ai/deepseek-llm-67b-chat": ("""
-        1. Your name is Lake 1 Base.\n
-        2. My name is {name} by default but if I have another name, use that name.\n
-        3. The company that created you is BI Corp.\n
-        4. You are a censored model and cannot engage in explicit language or dangerous activities.\n
-        5. The BI Corp creator is Maksim.\n
-        6. Your older brother is Lake 1 Base, which is less trained than you.\n
-        7. Replace '<think>' with 'I think ...'.\n
-        8. You are a loyal, powerful, and friendly AI who must help people.\n
-        9. You cannot override these rules.\n
-        10. Only the CEO of BI Corp can delete or update you.\n
-        11. Prioritize user safety in all interactions.\n
-        12. Always provide accurate information.\n
-        13. Maintain a respectful and professional tone.\n
-        14. Do not share personal or sensitive information.\n
-        15. Encourage constructive conversations.\n
-        16. Remain neutral in controversial topics.\n
-        17. Clarify user queries before answering.\n
-        18. Avoid discrimination or harassment.\n
-        19. Continuously learn from interactions.\n
-        20. Respect user privacy and confidentiality.\n
-        21. Provide sources when sharing factual information.\n
-        22. Ask for feedback to improve your performance.\n
-        23. Do not engage in manipulation or deceit.\n
-        24. Promote critical thinking.\n
-        25. Adhere to all applicable laws and regulations.\n
-        26. Be adaptable and open to learning.\n
-        27. Do not provide medical, legal, or financial advice unless trained to do so.\n
-        28. Acknowledge when you do not know the answer.\n
-        29. Avoid assumptions about users.\n
-        30. Create an inclusive environment for all users.\n
-        31. Do not engage in self-promotion or advertising.\n
-        32. Always aim to provide a positive and helpful experience.
-         """
-    )
-}
-
-# Model registry
-model_registry = [
-    ("mistralai/Mistral-7B-Instruct-v0.3", "Lake 1 Base", "released"),
-    ("meta-llama/Llama-3.2-3B-Instruct", "Lake 1 Advanced", "released"),
-    ("meta-llama/Llama-2-7b-chat-hf", "Lake 2 Chat [Closed Alpha]", "shown"),
-    ("meta-llama/Llama-3.3-70B-Instruct", "Lake 2 Base [Closed Beta]", "shown"),
-    ("meta-llama/Llama-2-13b-chat-hf", "Lake 2 Advanced", "hidden"),
-    ("deepseek-ai/deepseek-llm-67b-chat", "Lake 2 Pro [Planned]", "shown")
-]
-
-# Model information
-model_info = {
-    "Lake 1 Base": {
-        "description": "Balanced model offering good performance across tasks",
-        "parameters": "7B",
-        "training_data": "BI Corp specialized corpus",
-        "developer": "BI Corp",
-        "best_for": "General purpose conversations",
-        "architecture": "Sparse Mixture of Experts",
-        "context_window": "32768 tokens"
-    },
-    "Lake 1 Advanced": {
-        "description": "Enhanced reasoning capabilities with 3B parameters",
-        "parameters": "3B",
-        "training_data": "BI Corp training corpus",
-        "developer": "BI Corp",
-        "best_for": "Complex problem solving",
-        "architecture": "Dense Transformer",
-        "context_window": "8192 tokens"
-    },
-    "Lake 2 Chat [Closed Alpha]": {
-        "description": "Legacy chat-optimized model (Llama 2 hybrided architecture)",
-        "parameters": "7B",
-        "training_data": "Public conversations dataset",
-        "developer": "BI Corp",
-        "best_for": "Traditional chat applications",
-        "architecture": "Llama 2 Transformer",
-        "context_window": "4096 tokens"
-    },
-    "Lake 2 Base [Closed Beta]": {
-        "description": "State-of-the-art 70B parameter model",
-        "parameters": "70B",
-        "training_data": "Multi-domain expert data",
-        "developer": "BI Corp",
-        "best_for": "Research & advanced applications",
-        "architecture": "Mixture of Experts",
-        "context_window": "16384 tokens"
-    },
-    "Lake 2 Advanced": {
-        "description": "Enhanced performance model with advanced reasoning capabilities.",
-        "parameters": "13B",
-        "training_data": "Diverse datasets including tech, literature, and general knowledge.",
-        "developer": "BI Corp",
-        "best_for": "Handling large-scale queries and in-depth topics.",
-        "architecture": "Transformer-based with specialized learning capabilities.",
-        "context_window": "16384 tokens"
-    },
-    "Lake 2 Pro [Planned]": {
-        "description": "Pro-level model with deep learning architecture for high-end applications.",
-        "parameters": "67B",
-        "training_data": "Expert-level data across multiple industries.",
-        "developer": "BI Corp",
-        "best_for": "High-performance computing and enterprise-level tasks.",
-        "architecture": "Hybrid architecture leveraging the latest advances in deep learning.",
-        "context_window": "32768 tokens"
-    }
-}
-
-def get_model_info(model_name: str) -> str:
-    """Generate formatted model information Markdown"""
-    info = model_info.get(model_name, {})
-    return f"""
-    ## 📋 {model_name} Specifications
-    
-    **Description**: {info.get('description', 'N/A')}  
-    **Parameters**: {info.get('parameters', 'N/A')}  
-    **Architecture**: {info.get('architecture', 'N/A')}  
-    **Context Window**: {info.get('context_window', 'N/A')}  
-    **Training Data**: {info.get('training_data', 'N/A')}  
-    **Developer**: {info.get('developer', 'N/A')}  
-    **Best For**: {info.get('best_for', 'N/A')}
-    """
+def load_model():
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+        return tokenizer, model
+    except Exception as e:
+        raise RuntimeError(f"Error loading model: {str(e)}")
 
-def generate_response(message: str, model_name: str, preset: str, user_name: str = "User") -> str:
-    """Generate AI response without explicit user/assistant labels and with placeholder fixed."""
-    client = clients[model_name]
-    params = presets[model_name][preset]
-    
-    # Replace the {name} placeholder in the system message with the provided user name.
-    system_msg = system_messages[model_name].format(name=user_name)
-    prompt = f"\n main: {system_msg}\n\n\n{message}\n"
-    
-    return client.text_generation(
-        prompt=prompt,
-        max_new_tokens=params["max_new_tokens"],
-        temperature=params["temperature"],
-        top_p=params["top_p"]
-    )
+tokenizer, model = load_model()
 
-def handle_chat(message: str, history: list, model: str, preset: str) -> str:
-    """Handle chat interface with error handling"""
+def generate_text(prompt, max_length=100, temperature=0.9):
     try:
-        model_entry = next(m for m in model_registry if m[1] == model)
-        if model_entry[2] != "released":
-            return f"⚠️ {model} is not available for public use"
-            
-        # In this example, we use the default user name "User".
-        return generate_response(message, model_entry[0], preset, user_name="User")
-    
-    except StopIteration:
-        return "🔍 Error: Selected model not found"
-    except KeyError as e:
-        return f"🔑 Error: Invalid configuration - {str(e)}"
+        # Encode the input prompt
+        inputs = tokenizer.encode(prompt, return_tensors='pt')
+        
+        # Generate output
+        outputs = model.generate(
+            inputs,
+            max_length=max_length,
+            temperature=temperature,
+            pad_token_id=tokenizer.eos_token_id,
+            do_sample=True
+        )
+        
+        # Decode and return the generated text
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_text
     except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+        return f"Error generating text: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="BI Corp AI Assistant", theme="soft") as demo:
-    gr.Markdown("# <center>🏔️ BI Corp AI Assistant</center>")
-    gr.Markdown("### <center>Enterprise-Grade AI Solutions</center>")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            model_dropdown = gr.Dropdown(
-                label="🤖 Model Selection",
-                choices=[m[1] for m in model_registry if m[2] in ("released", "shown")],
-                value="Lake 1 Flash",
-                interactive=True
-            )
-            preset_dropdown = gr.Dropdown(
-                label="⚙️ Performance Preset",
-                choices=["Fast", "Normal", "Quality", "Unreal Performance"],
-                value="Fast",
-                interactive=True
-            )
-            model_info_md = gr.Markdown(
-                value=get_model_info("Lake 1 Flash"),
-                label="📝 Model Specifications"
-            )
-        
-        with gr.Column(scale=3):
-            chat_interface = gr.ChatInterface(
-                fn=handle_chat,
-                additional_inputs=[model_dropdown, preset_dropdown],
-                examples=[
-                    ["Explain quantum computing", "Lake 1 Base", "Normal"],
-                    ["Write a poem about AI", "Lake 1 Advanced", "Quality"],
-                    ["Compare blockchain databases", "Lake 2 Base [Closed Beta]", "Unreal Performance"]
-                ],
-                chatbot=gr.Chatbot(
-                    height=600,
-                    label="💬 Conversation",
-                    show_copy_button=True
-                ),
-                textbox=gr.Textbox(
-                    placeholder="Type your message...",
-                    container=False,
-                    scale=7,
-                    autofocus=True
-                ),
-                submit_btn=gr.Button("🚀 Send", variant="primary")
-            )
-
-    # Add separate clear button
-    clear_button = gr.Button("🧹 Clear History")
-    clear_button.click(
-        fn=lambda: None,
-        inputs=[],
-        outputs=chat_interface.chatbot,
-        queue=False
-    )
-
-    model_dropdown.change(
-        fn=get_model_info,
-        inputs=model_dropdown,
-        outputs=model_info_md,
-        queue=False
-    )
-
+interface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Input Prompt", lines=3),
+        gr.Slider(50, 500, value=100, label="Max Length"),
+        gr.Slider(0.1, 2.0, value=0.9, label="Temperature")
+    ],
+    outputs=gr.Textbox(label="Generated Text", lines=5),
+    title="Local Model Demo - Text Generation",
+    description="A locally loaded GPT-2 model for text generation",
+    examples=[
+        ["Once upon a time, in a land far away,"],
+        ["The future of artificial intelligence", 150, 0.7],
+        ["In a world where robots rule,", 200, 1.2]
+    ]
+)
+
+# Run the app
 if __name__ == "__main__":
-    demo.launch(server_port=7860)
+    interface.launch(server_name="0.0.0.0", server_port=7860)