Upload app.py with huggingface_hub

app.py

@@ -1,163 +1,79 @@
 """
-Myanmar LLM Gradio App
+Myanmar LLM Gradio App - Lite Version
 Model: amkyawdev/mm-llm-tiny
 """
 
-import os
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConfig
 import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Model name
 MODEL_NAME = "amkyawdev/mm-llm-tiny"
 
-# Load model and tokenizer
-print(f"Loading model: {MODEL_NAME}...")
-
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    
-    # Try with device_map first
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        low_cpu_mem_usage=True
-    )
-    
-    print("Model loaded on GPU!")
-    
-except Exception as e:
-    print(f"GPU failed: {e}, trying CPU...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
-    model = model.to("cpu")
-    print("Model loaded on CPU!")
+print(f"Loading {MODEL_NAME}...")
 
-# Create pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=256,
-    temperature=0.7,
-    top_p=0.95,
-)
+# Load tokenizer only first (saves memory)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token
 
-print("Model loaded!")
+# Model loads on first request (lazy load)
+model = None
 
-# System prompts
-SYSTEM_PROMPTS = {
-    "General": "သင်သည် မြန်မာစာကျွမ်းကျင်သော AI ဖြစ်ပါ။",
-    "Code Expert": "သင်သည် ပရိုဂရမ်ရေးသားတဲ့ ကျွမ်းကျင်သော AI ဖြစ်ပါ။",
-    "Translator": "သင်သည် ဘာသာပြန်ကျွမ်းကျင်သော AI ဖြစ်ပါ။",
-}
+def get_model():
+    global model
+    if model is None:
+        print("Loading model...")
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            torch_dtype=torch.float32,
+            low_cpu_mem_usage=True
+        )
+        model.eval()
+        print("Model loaded!")
+    return model
 
-def generate_response(message, history, system_prompt, max_tokens, temperature, top_p):
-    """Generate response"""
-    
-    # Build prompt
-    prompt = f"System: {system_prompt}\n\n"
+def generate(prompt, max_tokens=128, temp=0.7):
+    m = get_model()
     
-    for user_msg, bot_msg in history:
-        prompt += f"User: {user_msg}\n\nAssistant: {bot_msg}\n\n"
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256)
     
-    prompt += f"User: {message}\n\nAssistant:"
-    
-    # Generate
-    output = pipe(
-        prompt,
-        max_new_tokens=int(max_tokens),
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=temperature > 0,
-    )
-    
-    response = output[0]["generated_text"]
-    # Remove prompt from response
-    response = response[len(prompt):].strip()
+    with torch.no_grad():
+        outputs = m.generate(
+            **inputs,
+            max_new_tokens=int(max_tokens),
+            temperature=temp,
+            do_sample=temp > 0,
+        )
     
-    return response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response[len(prompt):].strip()
 
-# Build UI
+# UI
 with gr.Blocks(title="Myanmar LLM") as app:
-    
-    gr.Markdown("# 🇲🇲 Myanmar LLM Chat")
+    gr.Markdown("# 🇲🇲 Myanmar LLM")
     gr.Markdown("Model: **amkyawdev/mm-llm-tiny**")
     
     with gr.Row():
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(height=500)
-            
-            with gr.Row():
-                msg = gr.Textbox(
-                    label="Message",
-                    placeholder="မေးခွန်းရေးသားပါ။...",
-                    lines=3
-                )
-            
-            with gr.Row():
-                submit = gr.Button("📤 ပို့ပါ။", variant="primary")
-                clear = gr.Button("🗑️ သန့်ရှင်းပါ။")
-            
-            gr.Examples(
-                examples=[
-                    ["မင်္ဂလာပ�။"],
-                    ["Python နဲ့ Fibonacci ရေးပါ။"],
-                    ["Hello ကို မြန်မာလို ပြန်ပါ။"],
-                ],
-                inputs=msg
-            )
-        
-        with gr.Column(scale=1):
-            gr.Markdown("### ⚙️ Settings")
-            
-            system_prompt = gr.Dropdown(
-                choices=list(SYSTEM_PROMPTS.keys()),
-                value="General",
-                label="System Prompt"
-            )
-            
-            max_tokens = gr.Slider(
-                minimum=50, maximum=512, value=256, step=10,
-                label="Max Tokens"
-            )
-            
-            temperature = gr.Slider(
-                minimum=0.1, maximum=1.5, value=0.7, step=0.1,
-                label="Temperature"
-            )
-            
-            top_p = gr.Slider(
-                minimum=0.5, maximum=1.0, value=0.95, step=0.05,
-                label="Top-p"
-            )
+        msg = gr.Textbox(label="Message", placeholder="မေးခွန်းရေးသားပါ။...")
+        output = gr.Textbox(label="Response")
     
-    def respond(message, history, system_prompt, max_tokens, temperature, top_p):
-        response = generate_response(
-            message, history, 
-            SYSTEM_PROMPTS[system_prompt],
-            max_tokens, temperature, top_p
-        )
-        history.append((message, response))
-        return "", history
+    with gr.Row():
+        max_tokens = gr.Slider(32, 256, value=128, step=16, label="Max Tokens")
+        temp = gr.Slider(0.1, 1.0, value=0.7, label="Temperature")
     
-    submit.click(
-        respond,
-        inputs=[msg, chatbot, system_prompt, max_tokens, temperature, top_p],
-        outputs=[msg, chatbot]
-    )
+    btn = gr.Button("Generate")
     
-    msg.submit(
-        respond,
-        inputs=[msg, chatbot, system_prompt, max_tokens, temperature, top_p],
-        outputs=[msg, chatbot]
+    btn.click(
+        generate,
+        inputs=[msg, max_tokens, temp],
+        outputs=output
     )
     
-    clear.click(lambda: (None, []), outputs=[msg, chatbot])
+    gr.Examples(
+        examples=[
+            ["Hello မြန်မာလို ပြန်ပါ။", 64, 0.7],
+            ["Python နဲ့ list ရေးပါ။", 128, 0.7],
+        ],
+        inputs=[msg, max_tokens, temp]
+    )
 
 app.launch(share=True)