Update app.py

app.py

@@ -1,19 +1,19 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
 
-# Load a more capable open-source model
-strategy_generator = pipeline("text2text-generation", model="google/flan-t5-large")
+# Load a smaller, optimized model
+model_name = "google/flan-t5-base"  # Switch to a smaller model for faster inference
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Function to clean up repetitive or irrelevant sentences in generated text
-def clean_output(text):
-    sentences = text.split(". ")
-    seen = set()
-    cleaned_sentences = []
-    for sentence in sentences:
-        if sentence.strip() not in seen:
-            seen.add(sentence.strip())
-            cleaned_sentences.append(sentence.strip())
-    return ". ".join(cleaned_sentences)
+# Load model onto CPU with optimization
+strategy_generator = pipeline(
+    "text2text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1,  # Use GPU if available
+)
 
 # Function to generate actionable steps
 def generate_steps(industry, challenge, goals):
@@ -25,13 +25,12 @@ def generate_steps(industry, challenge, goals):
     Focus on specific, realistic, and innovative strategies relevant to the industry.
     """
     try:
-        response = strategy_generator(prompt, max_length=300, num_return_sequences=1, temperature=0.7, top_p=0.9)
-        cleaned_response = clean_output(response[0]['generated_text'])
-        return cleaned_response
+        response = strategy_generator(prompt, max_length=200, num_return_sequences=1, temperature=0.7, top_p=0.9)
+        return response[0]['generated_text']
     except Exception as e:
         return f"Error generating steps: {e}"
 
-# Function to generate "why" and "how" for each step
+# Function to combine rationale ("why") and implementation ("how")
 def expand_step(step):
     prompt = f"""
     You are a business consultant. For the following strategy:
@@ -41,25 +40,25 @@ def expand_step(step):
     - How to implement this step effectively.
     """
     try:
-        response = strategy_generator(prompt, max_length=200, num_return_sequences=1, temperature=0.7, top_p=0.9)
-        cleaned_response = clean_output(response[0]['generated_text'])
-        return cleaned_response
+        response = strategy_generator(prompt, max_length=150, num_return_sequences=1, temperature=0.7, top_p=0.9)
+        return response[0]['generated_text']
     except Exception as e:
         return f"Error expanding step: {e}"
 
-# Combined function to generate strategy
+# Combined function to generate detailed strategy
 def generate_strategy(industry, challenge, goals):
     # Generate initial steps
     steps = generate_steps(industry, challenge, goals)
     if "Error" in steps:
         return steps
     
-    # Expand each step
+    # Split steps and expand each
     steps_list = steps.split("\n")
     detailed_steps = []
     for step in steps_list:
         if step.strip():
-            detailed_steps.append(f"{step}\n{expand_step(step)}")
+            expanded = expand_step(step)
+            detailed_steps.append(f"{step}\n{expanded}")
     
     return "\n\n".join(detailed_steps)