Update app.py

app.py

@@ -1,67 +1,110 @@
 import gradio as gr
 from fastapi import FastAPI, Query
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-import uvicorn
+from functools import lru_cache
 import torch
+import logging
+
+# Setup logging untuk debugging performa
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Inisialisasi FastAPI
 app = FastAPI()
 
-# Inisialisasi model dan tokenizer
+# Preload model dan tokenizer untuk efisiensi
+logger.info("Memuat model saat startup...")
+
+# Cache model di memori
+model_cache = {}
+
 def load_model(model_name):
+    if model_name in model_cache:
+        logger.info(f"Menggunakan model {model_name} dari cache")
+        return model_cache[model_name]
+    
+    logger.info(f"Memuat model {model_name}...")
     if model_name == "mixtral":
         model_id = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)
-        return pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            # Gunakan 4-bit quantization untuk mengurangi penggunaan memori
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                load_in_4bit=True,  # Quantization untuk kecepatan
+                low_cpu_mem_usage=True
+            )
+            pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto")
+        except Exception as e:
+            logger.error(f"Gagal memuat Mixtral: {str(e)}")
+            raise
     elif model_name == "gpt2":
-        return pipeline("text-generation", model="gpt2")
+        pipe = pipeline("text-generation", model="gpt2", device=0 if torch.cuda.is_available() else -1)
     else:
         raise ValueError("Model tidak didukung. Pilih 'mixtral' atau 'gpt2'.")
+    
+    model_cache[model_name] = pipe
+    logger.info(f"Model {model_name} berhasil dimuat")
+    return pipe
+
+# Preload model saat startup
+try:
+    load_model("gpt2")  # Load GPT-2 (ringan) terlebih dahulu
+    load_model("mixtral")  # Load Mixtral dengan quantization
+except Exception as e:
+    logger.error(f"Error saat preload model: {str(e)}")
 
-# Fungsi untuk menghasilkan teks
-def generate_text(prompt, model_name, max_length=100):
+# Fungsi generate dengan caching
+@lru_cache(maxsize=100)
+def generate_text(prompt: str, model_name: str, max_length: int = 100):
     try:
-        generator = load_model(model_name)
-        # Menghasilkan teks
-        output = generator(prompt, max_length=max_length, num_return_sequences=1, do_sample=True)
+        logger.info(f"Memproses prompt: {prompt[:30]}... dengan model {model_name}")
+        generator = model_cache.get(model_name)
+        if not generator:
+            generator = load_model(model_name)
+        # Generate teks
+        output = generator(
+            prompt,
+            max_length=max_length,
+            num_return_sequences=1,
+            do_sample=True,
+            pad_token_id=generator.tokenizer.eos_token_id
+        )
         return output[0]["generated_text"]
     except Exception as e:
+        logger.error(f"Error saat generasi: {str(e)}")
         return f"Error: {str(e)}"
 
 # Endpoint API
 @app.get("/generate")
 async def generate(prompt: str = Query(..., description="Teks input untuk model"),
-                  model: str = Query("gpt2", description="Model AI: 'mixtral' atau 'gpt2'")):
-    result = generate_text(prompt, model)
+                  model: str = Query("gpt2", description="Model AI: 'mixtral' atau 'gpt2'"),
+                  max_length: int = Query(100, description="Panjang maksimum teks yang dihasilkan")):
+    result = generate_text(prompt, model, max_length)
     return {"prompt": prompt, "model": model, "generated_text": result}
 
 # Antarmuka Gradio
-def gradio_generate(prompt, model_choice):
-    return generate_text(prompt, model_choice)
+def gradio_generate(prompt, model_choice, max_length):
+    return generate_text(prompt, model_choice, max_length)
 
 with gr.Blocks() as demo:
-    gr.Markdown("# AI Text Generation API")
-    gr.Markdown("Masukkan teks dan pilih model untuk menghasilkan teks. Gunakan API di `/generate` untuk akses programatik.")
+    gr.Markdown("# AI Text Generation API (Optimized)")
+    gr.Markdown("Masukkan teks dan pilih model untuk menghasilkan teks. API tersedia di `/generate`.")
     
-    # Komponen input
     prompt_input = gr.Textbox(label="Prompt", placeholder="Masukkan teks di sini...")
     model_choice = gr.Dropdown(choices=["gpt2", "mixtral"], label="Pilih Model", value="gpt2")
+    max_length = gr.Slider(minimum=50, maximum=500, value=100, step=10, label="Panjang Maksimum")
     submit_button = gr.Button("Generate")
     
-    # Komponen output
     output_text = gr.Textbox(label="Hasil Generasi")
     
-    # Menghubungkan tombol dengan fungsi
     submit_button.click(
         fn=gradio_generate,
-        inputs=[prompt_input, model_choice],
+        inputs=[prompt_input, model_choice, max_length],
         outputs=output_text
     )
 
-# Menjalankan aplikasi (untuk lokal, bukan di Hugging Face)
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-else:
-    # Untuk Hugging Face Spaces, luncurkan Gradio
-    demo.launch()
+# Untuk Hugging Face Spaces, langsung launch Gradio
+demo.launch()