import gradio as gr
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
from peft import PeftModel

# 1. Define your models
# CHANGE THIS to the exact base model you used before fine-tuning (e.g., "t5-3b")
base_model_id = "google/flan-t5-xl" 

# Your fine-tuned repository
adapter_repo = "Spark2scale/Spark2Scale"

# 2. Load the Tokenizer
# We load this from your adapter repo to keep any special tokens you might have added
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    adapter_repo, 
    subfolder="checkpoint-4000" # Add this if your files are inside a specific folder
)

# 3. Load the Base Model in bfloat16
print("Loading base model in bfloat16...")
base_model = T5ForConditionalGeneration.from_pretrained(
    base_model_id,
    torch_dtype=torch.bfloat16
)

# 4. Apply your PEFT Adapter on top of the base model
print("Applying adapter weights...")
model = PeftModel.from_pretrained(
    base_model, 
    adapter_repo,
    subfolder="checkpoint-4000" # Remove this line if the files are in the main repo root
)
model.eval()

# 5. Define the prediction function
def evaluate_idea(startup_idea):
    inputs = tokenizer(startup_idea, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True
        )
        
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

# 6. Build and launch the Gradio Interface
iface = gr.Interface(
    fn=evaluate_idea,
    inputs=gr.Textbox(lines=5, placeholder="Enter the startup idea here...", label="Startup Idea"),
    outputs=gr.Textbox(label="Evaluation Report"),
    title="Spark2Scale Idea Evaluator",
    description="Backend API for the Spark2Scale evaluation agent."
)

iface.launch()