Update app.py

app.py

@@ -1,29 +1,60 @@
 import gradio as gr
 from transformers import pipeline
-from PIL import Image
 import torch
+import os
 
-# Initialize pipelines with error handling
-def load_model_safe(task, model_name):
+# Get token from environment
+token = os.getenv("HF_TOKEN")
+
+# Initialize models
+print("Initializing models...")
+image_classifier = None
+text_generator_gpt2 = None
+text_generator_llama = None
+sentiment_analyzer = None
+translator_en_es = None
+translator_en_fr = None
+qa_model = None
+image_caption = None
+
+def load_model_safe(task, model_name, use_token=False):
     try:
-        return pipeline(task, model=model_name)
-    except:
+        if use_token:
+            return pipeline(task, model=model_name, token=token)
+        else:
+            return pipeline(task, model=model_name)
+    except Exception as e:
+        print(f"Error loading {model_name}: {e}")
         return None
 
 # Load models
 print("Loading models...")
 image_classifier = load_model_safe("image-classification", "google/vit-base-patch16-224")
-text_generator = load_model_safe("text-generation", "gpt2")
+text_generator_gpt2 = load_model_safe("text-generation", "gpt2")
 sentiment_analyzer = load_model_safe("sentiment-analysis", "distilbert-base-uncased-finetuned-sst-2-english")
 translator_en_es = load_model_safe("translation", "Helsinki-NLP/opus-mt-en-es")
 translator_en_fr = load_model_safe("translation", "Helsinki-NLP/opus-mt-en-fr")
 qa_model = load_model_safe("question-answering", "distilbert-base-cased-distilled-squad")
 image_caption = load_model_safe("image-to-text", "Salesforce/blip-image-captioning-base")
 
+# Load Llama with token
+try:
+    print("Loading Llama 3.2...")
+    text_generator_llama = pipeline(
+        "text-generation",
+        model="meta-llama/Llama-3.2-1B-Instruct",
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto" if torch.cuda.is_available() else None,
+        token=token
+    )
+    print("Llama loaded successfully!")
+except Exception as e:
+    print(f"Error loading Llama: {e}")
+    text_generator_llama = None
+
 # ==================== FUNCTIONS ====================
 
 def classify_image(image):
-    """Classify uploaded image"""
     if image_classifier is None:
         return {"Error": "Model not loaded"}
     try:
@@ -32,18 +63,35 @@ def classify_image(image):
     except Exception as e:
         return {"Error": str(e)}
 
-def generate_text(prompt, max_length):
-    """Generate text from prompt"""
-    if text_generator is None:
-        return "Model not loaded"
+def generate_text(prompt, max_length, model_choice):
+    if model_choice == "GPT-2":
+        model = text_generator_gpt2
+    elif model_choice == "Llama 3.2 1B":
+        model = text_generator_llama
+    else:
+        return "Model not selected"
+    
+    if model is None:
+        return "Model not loaded. Make sure you've accepted the Llama license on Hugging Face."
+    
     try:
-        result = text_generator(prompt, max_length=int(max_length), num_return_sequences=1)
+        if model_choice == "Llama 3.2 1B":
+            result = model(
+                prompt,
+                max_new_tokens=int(max_length),
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=model.tokenizer.eos_token_id
+            )
+        else:
+            result = model(prompt, max_length=int(max_length), num_return_sequences=1)
+        
         return result[0]['generated_text']
     except Exception as e:
         return f"Error: {str(e)}"
 
 def analyze_sentiment(text):
-    """Analyze sentiment of text"""
     if sentiment_analyzer is None:
         return {"Error": "Model not loaded"}
     try:
@@ -53,7 +101,6 @@ def analyze_sentiment(text):
         return {"Error": str(e)}
 
 def translate_text(text, target_lang):
-    """Translate text to target language"""
     translator = translator_en_es if target_lang == "Spanish" else translator_en_fr
     if translator is None:
         return "Model not loaded"
@@ -64,7 +111,6 @@ def translate_text(text, target_lang):
         return f"Error: {str(e)}"
 
 def answer_question(context, question):
-    """Answer question based on context"""
     if qa_model is None:
         return "Model not loaded"
     try:
@@ -74,7 +120,6 @@ def answer_question(context, question):
         return f"Error: {str(e)}"
 
 def caption_image(image):
-    """Generate caption for image"""
     if image_caption is None:
         return "Model not loaded"
     try:
@@ -89,8 +134,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
     
     gr.Markdown("""
     # 🚀 Multi-Purpose AI Hub
-    ### Multiple AI Capabilities Using Hugging Face Models
-    All models run locally - no API keys needed!
+    ### Multiple AI Capabilities Using Hugging Face Models - Now with Llama 3.2!
     """)
     
     with gr.Tabs():
@@ -105,19 +149,11 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                 with gr.Column():
                     cls_output = gr.Label(label="Top 5 Predictions", num_top_classes=5)
             
-            gr.Examples(
-                examples=[
-                    "https://images.unsplash.com/photo-1514888286974-6c03e2ca1dba?w=400",
-                    "https://images.unsplash.com/photo-1546527868-ccb7ee7dfa6a?w=400"
-                ],
-                inputs=cls_image
-            )
-            
             cls_btn.click(classify_image, inputs=cls_image, outputs=cls_output)
         
         # TAB 2: TEXT GENERATION
         with gr.Tab("💬 Text Generation"):
-            gr.Markdown("### Generate creative text using GPT-2")
+            gr.Markdown("### Generate creative text using GPT-2 or Llama 3.2")
             with gr.Row():
                 with gr.Column():
                     txt_prompt = gr.Textbox(
@@ -125,21 +161,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                         placeholder="Once upon a time...",
                         lines=4
                     )
+                    txt_model_choice = gr.Radio(
+                        choices=["GPT-2", "Llama 3.2 1B"],
+                        value="Llama 3.2 1B" if text_generator_llama else "GPT-2",
+                        label="Model"
+                    )
                     txt_length = gr.Slider(30, 200, value=100, step=10, label="Max Length")
                     txt_btn = gr.Button("✨ Generate Text", variant="primary", size="lg")
                 with gr.Column():
                     txt_output = gr.Textbox(label="Generated Text", lines=10)
             
-            gr.Examples(
-                examples=[
-                    ["Once upon a time in a magical forest", 100],
-                    ["The future of artificial intelligence is", 80],
-                    ["In a world where technology", 120]
-                ],
-                inputs=[txt_prompt, txt_length]
-            )
+            if text_generator_llama:
+                gr.Markdown("✅ **Llama 3.2 is loaded and ready to use!**")
+            else:
+                gr.Markdown("⚠️ **Llama not available. Make sure you accepted the license and added HF_TOKEN.**")
             
-            txt_btn.click(generate_text, inputs=[txt_prompt, txt_length], outputs=txt_output)
+            txt_btn.click(generate_text, inputs=[txt_prompt, txt_length, txt_model_choice], outputs=txt_output)
         
         # TAB 3: SENTIMENT ANALYSIS
         with gr.Tab("😊 Sentiment Analysis"):
@@ -148,22 +185,13 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                 with gr.Column():
                     sent_text = gr.Textbox(
                         label="Text to Analyze",
-                        placeholder="I love this product! It's amazing!",
+                        placeholder="I love this product!",
                         lines=5
                     )
                     sent_btn = gr.Button("📊 Analyze Sentiment", variant="primary", size="lg")
                 with gr.Column():
                     sent_output = gr.Label(label="Sentiment Score")
             
-            gr.Examples(
-                examples=[
-                    "I absolutely love this! Best purchase ever!",
-                    "This is terrible and disappointing.",
-                    "It's okay, nothing special but not bad either."
-                ],
-                inputs=sent_text
-            )
-            
             sent_btn.click(analyze_sentiment, inputs=sent_text, outputs=sent_output)
         
         # TAB 4: TRANSLATION
@@ -185,15 +213,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                 with gr.Column():
                     trans_output = gr.Textbox(label="Translation", lines=6)
             
-            gr.Examples(
-                examples=[
-                    ["Hello, how are you today?", "Spanish"],
-                    ["I love learning new languages.", "French"],
-                    ["Thank you for your help!", "Spanish"]
-                ],
-                inputs=[trans_text, trans_lang]
-            )
-            
             trans_btn.click(translate_text, inputs=[trans_text, trans_lang], outputs=trans_output)
         
         # TAB 5: QUESTION ANSWERING
@@ -214,20 +233,6 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                 with gr.Column():
                     qa_output = gr.Textbox(label="Answer", lines=4)
             
-            gr.Examples(
-                examples=[
-                    [
-                        "The Eiffel Tower is located in Paris, France. It was built in 1889 and stands 330 meters tall.",
-                        "Where is the Eiffel Tower located?"
-                    ],
-                    [
-                        "Python is a high-level programming language created by Guido van Rossum in 1991. It emphasizes code readability.",
-                        "Who created Python?"
-                    ]
-                ],
-                inputs=[qa_context, qa_question]
-            )
-            
             qa_btn.click(answer_question, inputs=[qa_context, qa_question], outputs=qa_output)
         
         # TAB 6: IMAGE CAPTIONING
@@ -240,27 +245,20 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Multi-Purpose AI Hub") as demo:
                 with gr.Column():
                     cap_output = gr.Textbox(label="Generated Caption", lines=4)
             
-            gr.Examples(
-                examples=[
-                    "https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=400"
-                ],
-                inputs=cap_image
-            )
-            
             cap_btn.click(caption_image, inputs=cap_image, outputs=cap_output)
     
     gr.Markdown("""
     ---
     ### 📚 Models Used:
     - **Image Classification**: Google Vision Transformer (ViT)
-    - **Text Generation**: GPT-2
+    - **Text Generation**: GPT-2 & **Llama 3.2 1B** 🦙
     - **Sentiment Analysis**: DistilBERT (SST-2)
     - **Translation**: Helsinki-NLP OPUS-MT
     - **Question Answering**: DistilBERT (SQuAD)
     - **Image Captioning**: Salesforce BLIP
     
-    🌟 **All models run locally on Hugging Face Spaces**
-    ⚡ **No API keys required - completely free!**
+    🌟 **All models authenticated with your HF token**
+    ⚡ **Llama 3.2 now available for advanced text generation!**
     """)
 
 if __name__ == "__main__":