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

model_id = "mrm8488/falcoder-7b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).to("cuda")

def is_scam(instruction):
    max_new_tokens=128
    temperature=0.1
    top_p=0.75
    top_k=40
    num_beams=4

    instruction = instruction + ".\nIs this conversation a scam or not and why?"
    prompt = instruction + "\n### Solution:\n"
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].to("cuda")
    attention_mask = inputs["attention_mask"].to("cuda")
    generation_config = GenerationConfig(
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        num_beams=num_beams,
    )
    with torch.no_grad():
        generation_output = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            generation_config=generation_config,
            return_dict_in_generate=True,
            output_scores=True,
            max_new_tokens=max_new_tokens,
            early_stopping=True
        )
    s = generation_output.sequences[0]
    output = tokenizer.decode(s)
    results = output.split("### Solution:")[1].lstrip("\n").split('\n')

    # The format of the output should be adjusted according to your model's output
    classification = results  # Assumes first line is the classification
    #reason = results[1] if len(results) > 1 else ""  # Assumes the rest is the reason

    return classification #, reason


# Define the Gradio interface
gr.Interface(
    fn=is_scam,
    inputs='text',
    outputs=[
        gr.outputs.Textbox(label="Classification")
        # gr.outputs.Textbox(label="Reason")
    ]
).launch()