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

model_name = "mistralai/Mistral-7B-Instruct-v0.2"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

q_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)

def build_prompt(context, num_questions):
    return (
        f"You are an expert interview question generator. "
        f"Generate {num_questions} concise and relevant interview questions based on the following topic or paragraph:\n\n"
        f"{context.strip()}\n\nQuestions:"
    )

def generate_questions(context, num_questions):
    prompt = build_prompt(context, num_questions)
    output = q_pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_p=0.9)
    return output[0]['generated_text'].split("Questions:")[-1].strip()

iface = gr.Interface(
    fn=generate_questions,
    inputs=[
        gr.Textbox(lines=4, label="Enter a topic or paragraph"),
        gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Number of Questions")
    ],
    outputs="text",
    title="Mistral Interview Question Generator",
    description="Generates interview questions using the Mistral-7B-Instruct model in 4-bit for efficient memory usage."
)

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