Update Dockerfile

app.py

@@ -1,45 +1,109 @@
 
+# import gradio as gr
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+# from huggingface_hub import login
+# import torch
+# import os
+
+# # Authenticate using environment variable
+# login(token=os.getenv('HF_TOKEN'))
+
+# # Load model (will use cached version if available)
+# model_id = "meta-llama/Llama-2-7b-chat-hf"
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# def load_model():
+#     tokenizer = AutoTokenizer.from_pretrained(model_id)
+#     model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
+#     return tokenizer, model
+
+# tokenizer, model = load_model()
+
+# def generate_text(prompt, max_length=200):
+#     inputs = tokenizer(prompt, return_tensors="pt").to(device)
+#     outputs = model.generate(
+#         **inputs,
+#         max_new_tokens=max_length,
+#         temperature=0.7,
+#         do_sample=True
+#     )
+#     return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# # Gradio interface
+# with gr.Blocks() as demo:
+#     gr.Markdown("# LLaMA 2 7B Chat Demo")
+#     with gr.Row():
+#         input_text = gr.Textbox(label="Input Prompt", lines=3)
+#         output_text = gr.Textbox(label="Generated Response", lines=3)
+#     generate_btn = gr.Button("Generate")
+#     generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)
+
+# demo.launch(server_name="0.0.0.0", server_port=7860)
+
+
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login, hf_hub_download
+from tenacity import retry, stop_after_attempt, wait_exponential
 import torch
+import os
 
-# Load model (will use cached version if available)
-model_id = "meta-llama/Llama-2-7b-chat-hf"
+# Authentication
+login(token=os.getenv('HF_TOKEN'))
 
-# Check for GPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Configuration
+CACHE_REPO = "Juna190825/cacheRepo"  # Your dataset repo for cached models
+MODEL_ID = "meta-llama/Llama-2-7b-chat-hf"  # Original model ID
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
+@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
+def load_model():
+    try:
+        # First try loading from cache repo
+        model = AutoModelForCausalLM.from_pretrained(
+            CACHE_REPO,
+            cache_dir="/cache/models",
+            local_files_only=True
+        ).to(DEVICE)
+        tokenizer = AutoTokenizer.from_pretrained(
+            CACHE_REPO,
+            cache_dir="/cache/models"
+        )
+        print("Loaded model from cache repo")
+        return model, tokenizer
+    except Exception as e:
+        print(f"Cache load failed: {str(e)}. Falling back to original repo")
+        # Fallback to original repo
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            cache_dir="/cache/models"
+        ).to(DEVICE)
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            cache_dir="/cache/models"
+        )
+        return model, tokenizer
+
+# Load model and tokenizer
+model, tokenizer = load_model()
 
 def generate_text(prompt, max_length=200):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
-    # Generate response
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
     outputs = model.generate(
         **inputs,
         max_new_tokens=max_length,
         temperature=0.7,
         do_sample=True
     )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Create Gradio interface
+# Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# LLaMA 2 7B Chat Demo")
     with gr.Row():
         input_text = gr.Textbox(label="Input Prompt", lines=3)
         output_text = gr.Textbox(label="Generated Response", lines=3)
-    
     generate_btn = gr.Button("Generate")
-    generate_btn.click(
-        fn=generate_text,
-        inputs=input_text,
-        outputs=output_text
-    )
+    generate_btn.click(fn=generate_text, inputs=input_text, outputs=output_text)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)
-