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

MODEL_PATH = "."

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def resumebot(question):
    prompt = f"Answer this resume question: {question}"
    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=256
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=100,
            num_beams=4,
            early_stopping=True
        )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=resumebot,
    inputs=gr.Textbox(
        lines=2,
        placeholder="Ask something about Vaishnavi's resume..."
    ),
    outputs=gr.Textbox(label="ResumeBot Answer"),
    title="Vaishnavi ResumeBot",
    description="Ask questions about education, skills, projects and experience."
)

if __name__ == "__main__":
    demo.launch()