import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "IlmaJiyadh/phi3-4k-ft"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    trust_remote_code=True
)

def summarize(text):
    prompt = f"Below is a lecture transcript. Take lecture notes in bullet points.\n\nInput:\n{text}\n\nSummary:\n"
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7, use_cache=False)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

gr.Interface(
    fn=summarize,
    inputs=gr.Textbox(lines=10, label="📄 Paste Transcript"),
    outputs=gr.Textbox(label="📝 Summary"),
    title="🧠 Transcript → Summary (Phi-3 Fine-tuned)",
    description="Test only the summarization step."
).launch()