Hm

app.py

@@ -23,26 +23,55 @@ def load_model(hf_token):
     global global_model, global_tokenizer, model_loaded
     
     if not hf_token:
-        return False, "Please enter your Hugging Face token to use the model."
+        model_loaded = False
+        return "⚠️ Please enter your Hugging Face token to use the model."
     
-    model_name = "google/gemma-3-4b-pt"
     try:
-        global_tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-        global_model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            token=hf_token
-        )
-        model_loaded = True
-        return True, "Model loaded successfully!"
+        # Try both model versions
+        model_options = [
+            "google/gemma-3-4b-pt",    # Try the quantized PT version first
+            "google/gemma-2b",         # Fallback to 2b model
+        ]
+        
+        # Try to load models in order until one works
+        for model_name in model_options:
+            try:
+                print(f"Attempting to load model: {model_name}")
+                
+                # Load tokenizer
+                global_tokenizer = AutoTokenizer.from_pretrained(
+                    model_name, 
+                    token=hf_token
+                )
+                
+                # Load model with safe configuration
+                global_model = AutoModelForCausalLM.from_pretrained(
+                    model_name,
+                    torch_dtype=torch.float16,
+                    device_map="auto",
+                    token=hf_token,
+                    use_cache=True,
+                    low_cpu_mem_usage=True,
+                    attn_implementation="flash_attention_2" if torch.cuda.is_available() else "eager"
+                )
+                
+                model_loaded = True
+                return f"✅ Model {model_name} loaded successfully!"
+            except Exception as specific_e:
+                print(f"Failed to load {model_name}: {specific_e}")
+                continue
+                
+        # If we get here, all model options failed
+        model_loaded = False
+        return "❌ Could not load any model version. Please check your token and try again."
+        
     except Exception as e:
         model_loaded = False
         error_msg = str(e)
         if "401 Client Error" in error_msg:
-            return False, "Authentication failed. Please check your token and make sure you've accepted the model license on Hugging Face."
+            return "❌ Authentication failed. Please check your token and make sure you've accepted the model license on Hugging Face."
         else:
-            return False, f"Error loading model: {error_msg}"
+            return f"❌ Error loading model: {error_msg}"
 
 def generate_prompt(task_type, **kwargs):
     """Generate appropriate prompts based on task type and parameters"""
@@ -135,21 +164,37 @@ def generate_text(prompt, max_length=1024, temperature=0.7, top_p=0.95):
     try:
         inputs = global_tokenizer(prompt, return_tensors="pt").to(global_model.device)
         
-        # Generate text
+        generation_config = {
+            "max_length": max_length,
+            "do_sample": True,
+            "pad_token_id": global_tokenizer.eos_token_id,
+        }
+        
+        # Only add temperature if it's not too low (can cause probability issues)
+        if temperature >= 0.2:
+            generation_config["temperature"] = temperature
+        else:
+            generation_config["temperature"] = 0.2
+            
+        # Only add top_p if it's valid
+        if 0 < top_p < 1:
+            generation_config["top_p"] = top_p
+            
+        # Generate text with safer parameters
         outputs = global_model.generate(
             **inputs,
-            max_length=max_length,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=global_tokenizer.eos_token_id,
+            **generation_config
         )
         
         # Decode and return the generated text
         generated_text = global_tokenizer.decode(outputs[0], skip_special_tokens=True)
         return generated_text
     except Exception as e:
-        return f"Error generating text: {str(e)}"
+        error_msg = str(e)
+        if "probability tensor" in error_msg:
+            return "Error: There was a problem with the generation parameters. Try using higher temperature (0.5+) and top_p values (0.9+)."
+        else:
+            return f"Error generating text: {error_msg}"
 
 # Create parameters UI component
 def create_parameter_ui():
@@ -162,16 +207,16 @@ def create_parameter_ui():
             label="Maximum Length"
         )
         temperature = gr.Slider(
-            minimum=0.1, 
+            minimum=0.3, 
             maximum=1.5, 
-            value=0.7, 
+            value=0.8, 
             step=0.1,
             label="Temperature"
         )
         top_p = gr.Slider(
             minimum=0.5, 
-            maximum=1.0, 
-            value=0.95, 
+            maximum=0.99, 
+            value=0.9, 
             step=0.05,
             label="Top-p"
         )
@@ -330,10 +375,10 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for text generation
             gr.Examples(
                 [
-                    ["Creative Writing", "short story", "a robot learning to paint", "article", "artificial intelligence", "", 1024, 0.7, 0.95],
-                    ["Creative Writing", "poem", "the beauty of mathematics", "article", "artificial intelligence", "", 768, 0.8, 0.95],
-                    ["Informational Writing", "short story", "a robot discovering emotions", "article", "quantum computing", "", 1024, 0.5, 0.95],
-                    ["Custom Prompt", "short story", "a robot discovering emotions", "article", "artificial intelligence", "Write a marketing email for a new smartphone with innovative AI features", 1024, 0.7, 0.95]
+                    ["Creative Writing", "short story", "a robot learning to paint", "article", "artificial intelligence", "", 1024, 0.8, 0.9],
+                    ["Creative Writing", "poem", "the beauty of mathematics", "article", "artificial intelligence", "", 768, 0.8, 0.9],
+                    ["Informational Writing", "short story", "a robot discovering emotions", "article", "quantum computing", "", 1024, 0.7, 0.9],
+                    ["Custom Prompt", "short story", "a robot discovering emotions", "article", "artificial intelligence", "Write a marketing email for a new smartphone with innovative AI features", 1024, 0.8, 0.9]
                 ],
                 fn=text_generation_handler,
                 inputs=[
@@ -390,9 +435,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for brainstorming
             gr.Examples(
                 [
-                    ["project", "educational app for children", 1024, 0.8, 0.95],
-                    ["business", "eco-friendly food packaging", 1024, 0.8, 0.95],
-                    ["solution", "reducing urban traffic congestion", 1024, 0.8, 0.95],
+                    ["project", "educational app for children", 1024, 0.8, 0.9],
+                    ["business", "eco-friendly food packaging", 1024, 0.8, 0.9],
+                    ["solution", "reducing urban traffic congestion", 1024, 0.8, 0.9],
                 ],
                 fn=brainstorm_handler,
                 inputs=[brainstorm_category, brainstorm_topic, *brainstorm_params],
@@ -452,9 +497,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for content creation
             gr.Examples(
                 [
-                    ["blog post", "sustainable living tips", "environmentally conscious consumers", 1536, 0.7, 0.95],
-                    ["social media post", "product launch announcement", "existing customers", 512, 0.7, 0.95],
-                    ["marketing copy", "new fitness app", "health-focused individuals", 1024, 0.7, 0.95],
+                    ["blog post", "sustainable living tips", "environmentally conscious consumers", 1536, 0.8, 0.9],
+                    ["social media post", "product launch announcement", "existing customers", 512, 0.8, 0.9],
+                    ["marketing copy", "new fitness app", "health-focused individuals", 1024, 0.8, 0.9],
                 ],
                 fn=content_creation_handler,
                 inputs=[content_type, content_topic, content_audience, *content_params],
@@ -509,9 +554,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for email drafting
             gr.Examples(
                 [
-                    ["job application", "Applying for a marketing specialist position at ABC Marketing. I have 5 years of experience in digital marketing.", 1024, 0.7, 0.95],
-                    ["business proposal", "Proposing a collaboration between our companies for a joint product development effort.", 1024, 0.7, 0.95],
-                    ["follow-up", "Following up after our meeting last Thursday about the project timeline and resources.", 1024, 0.7, 0.95],
+                    ["job application", "Applying for a marketing specialist position at ABC Marketing. I have 5 years of experience in digital marketing.", 1024, 0.8, 0.9],
+                    ["business proposal", "Proposing a collaboration between our companies for a joint product development effort.", 1024, 0.8, 0.9],
+                    ["follow-up", "Following up after our meeting last Thursday about the project timeline and resources.", 1024, 0.8, 0.9],
                 ],
                 fn=email_draft_handler,
                 inputs=[email_type, email_context, *email_params],
@@ -609,9 +654,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for explanation
             gr.Examples(
                 [
-                    ["blockchain technology", "beginner", 1024, 0.7, 0.95],
-                    ["photosynthesis", "child", 1024, 0.7, 0.95],
-                    ["machine learning", "college student", 1024, 0.7, 0.95],
+                    ["blockchain technology", "beginner", 1024, 0.8, 0.9],
+                    ["photosynthesis", "child", 1024, 0.8, 0.9],
+                    ["machine learning", "college student", 1024, 0.8, 0.9],
                 ],
                 fn=explanation_handler,
                 inputs=[explain_topic, explain_level, *explain_params],
@@ -666,9 +711,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for classification
             gr.Examples(
                 [
-                    ["The stock market saw significant gains today as tech companies reported strong quarterly earnings.", "technology, health, finance, entertainment, education, sports", 256, 0.1, 0.95],
-                    ["The team scored in the final minutes to secure their victory in the championship game.", "technology, health, finance, entertainment, education, sports", 256, 0.1, 0.95],
-                    ["The new educational app helps students master complex math concepts through interactive exercises.", "technology, health, finance, entertainment, education, sports", 256, 0.1, 0.95],
+                    ["The stock market saw significant gains today as tech companies reported strong quarterly earnings.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
+                    ["The team scored in the final minutes to secure their victory in the championship game.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
+                    ["The new educational app helps students master complex math concepts through interactive exercises.", "technology, health, finance, entertainment, education, sports", 256, 0.5, 0.9],
                 ],
                 fn=classification_handler,
                 inputs=[classify_text, classify_categories, *classify_params],
@@ -723,8 +768,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             # Examples for data extraction
             gr.Examples(
                 [
-                    ["Sarah Johnson is the CEO of Green Innovations, founded in 2012. The company reported $8.5 million in revenue for 2023. Contact her at [email protected].", "name, position, company, founding year, revenue, contact", 768, 0.3, 0.95],
-                    ["The new iPhone 15 Pro features a 6.1-inch display, A17 Pro chip, 48MP camera, and starts at $999 for the 128GB model.", "product name, screen size, processor, camera, price, storage capacity", 768, 0.3, 0.95],
+                    ["Sarah Johnson is the CEO of Green Innovations, founded in 2012. The company reported $8.5 million in revenue for 2023. Contact her at [email protected].", "name, position, company, founding year, revenue, contact", 768, 0.5, 0.9],
+                    ["The new iPhone 15 Pro features a 6.1-inch display, A17 Pro chip, 48MP camera, and starts at $999 for the 128GB model.", "product name, screen size, processor, camera, price, storage capacity", 768, 0.5, 0.9],
                 ],
                 fn=data_extraction_handler,
                 inputs=[extract_text, extract_data_points, *extract_params],