Update app.py

app.py

@@ -1,173 +1,56 @@
-import gc
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
-import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, login
+import torch
 
-# Load Hugging Face token
 HF_TOKEN = os.getenv("HF_TOKEN")
 login(token=HF_TOKEN)
 
-# Define models
-MODELS = {
-    "atlas-flash-1215": {
-        "name": "🦁 Atlas-Flash 1215",
-        "sizes": {
-            "1.5B": "Spestly/Atlas-Flash-1.5B-Preview",
-        },
-        "emoji": "🦁",
-        "experimental": True,
-        "is_vision": False,
-        "system_prompt_env": "ATLAS_FLASH_1215",
-    },
-    "atlas-pro-0403": {
-        "name": "🏆 Atlas-Pro 0403",
-        "sizes": {
-            "1.5B": "Spestly/Atlas-Pro-1.5B-Preview",
-        },
-        "emoji": "🏆",
-        "experimental": True,
-        "is_vision": False,
-        "system_prompt_env": "ATLAS_PRO_0403",
-    },
-}
-
-# Clear memory
-def clear_memory():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    gc.collect()
-
-# Load model
-def load_model(model_key, model_size):
-    try:
-        clear_memory()
-        
-        # Unload previous model if any
-        global current_model
-        if current_model is not None:
-            del current_model["model"]
-            del current_model["tokenizer"]
-            clear_memory()
-
-        model_path = MODELS[model_key]["sizes"][model_size]
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            trust_remote_code=True,
-            low_cpu_mem_usage=True
-        )
-
-        current_model.update({
-            "tokenizer": tokenizer,
-            "model": model,
-            "config": {
-                "name": f"{MODELS[model_key]['name']} {model_size}",
-                "path": model_path,
-                "system_prompt": os.getenv(MODELS[model_key]["system_prompt_env"], "Default system prompt"),
-            }
-        })
-        return f"✅ {MODELS[model_key]['name']} {model_size} loaded successfully!"
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
-
-# Respond to input
-def respond(prompt, max_tokens, temperature, top_p, top_k):
-    if not current_model["model"] or not current_model["tokenizer"]:
-        return "⚠️ Please select and load a model first"
-
-    try:
-        system_prompt = current_model["config"]["system_prompt"]
-        if not system_prompt:
-            return "⚠️ System prompt not found for the selected model."
-
-        full_prompt = f"{system_prompt}\n\n### Instruction:\n{prompt}\n\n### Response:"
-
-        inputs = current_model["tokenizer"](
-            full_prompt,
-            return_tensors="pt",
-            max_length=512,
-            truncation=True,
-            padding=True
-        )
-        with torch.no_grad():
-            output = current_model["model"].generate(
-                input_ids=inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=max_tokens,
-                temperature=temperature,
-                top_p=top_p,
-                top_k=top_k,
-                do_sample=True,
-                pad_token_id=current_model["tokenizer"].pad_token_id,
-                eos_token_id=current_model["tokenizer"].eos_token_id,
-            )
-            response = current_model["tokenizer"].decode(output[0], skip_special_tokens=True)
 
-        if full_prompt in response:
-            response = response.replace(full_prompt, "").strip()
+model_name = "Spestly/Atlas-Pro-1.5B-Preview"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True)
 
-        return response
-    except Exception as e:
-        return f"⚠️ Generation Error: {str(e)}"
-    finally:
-        clear_memory()
 
-# Initialize model storage
-current_model = {"tokenizer": None, "model": None, "config": None}
+model.eval()
 
-# UI for Gradio
-def gradio_ui():
-    def load_and_set_model(model_key, model_size):
-        return load_model(model_key, model_size)
+def generate_response(message, history):
+    instruction = (
+        "You are an LLM called Atlas. You are finetuned by Aayan Mishra. You are NOT trained by Anthropic. "
+        "You are a Qwen 2.5 fine-tune. Your purpose is the help the user accomplish their request to the best of your abilities. "
+        "Below is an instruction that describes a task. Answer it clearly and concisely.\n\n"
+        f"### Instruction:\n{message}\n\n### Response:"
+    )
     
-    with gr.Blocks() as app:
-        gr.Markdown("## 🦁 Atlas Inference Platform - Experimental 🧪")
-        
-        with gr.Row():
-            model_key_dropdown = gr.Dropdown(
-                choices=list(MODELS.keys()),
-                value=list(MODELS.keys())[0],
-                label="Select Model Variant",
-                interactive=True
-            )
-            model_size_dropdown = gr.Dropdown(
-                choices=list(MODELS[list(MODELS.keys())[0]]["sizes"].keys()),
-                value="1.5B",
-                label="Select Model Size",
-                interactive=True
-            )
-            load_button = gr.Button("Load Model")
-
-        load_status = gr.Textbox(label="Model Load Status", interactive=False)
-
-        load_button.click(
-            load_and_set_model,
-            inputs=[model_key_dropdown, model_size_dropdown],
-            outputs=load_status,
-        )
-
-        with gr.Row():
-            prompt_input = gr.Textbox(label="Input Prompt", lines=4)
-            max_tokens_slider = gr.Slider(10, 512, value=256, step=10, label="Max Tokens")
-            temperature_slider = gr.Slider(0.1, 2.0, value=0.4, step=0.1, label="Temperature")
-            top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="Top-P")
-            top_k_slider = gr.Slider(1, 100, value=50, step=1, label="Top-K")
-        
-        generate_button = gr.Button("Generate Response")
-        response_output = gr.Textbox(label="Model Response", lines=6, interactive=False)
-
-        generate_button.click(
-            respond,
-            inputs=[prompt_input, max_tokens_slider, temperature_slider, top_p_slider, top_k_slider],
-            outputs=response_output,
+    inputs = tokenizer(instruction, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=1000,
+            num_return_sequences=1,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
         )
-
-    return app
-
-if __name__ == "__main__":
-    gradio_ui().launch()
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.split("### Response:")[-1].strip()
+    
+    return response
+
+iface = gr.ChatInterface(
+    generate_response,
+    chatbot=gr.Chatbot(height=600, type="messages"),
+    textbox=gr.Textbox(placeholder="Type your message here...", container=False, scale=7),
+    title="🦁 Atlas-Pro",
+    description="Chat with Alas-Pro",
+    theme="citrus",
+    examples=[
+        "Can you give me a good salsa recipe?",
+        "Write an engaging two-line horror story.",
+        "What is the capital of Australia?",
+    ],
+    type="messages"
+)
+
+iface.launch()