import gradio as gr
from transformers import AutoTokenizer
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM

# Load the PEFT model configuration and model
config = PeftConfig.from_pretrained("isashap/contexttrained-validationloss-gpt2FINAL")
base_model = AutoModelForCausalLM.from_pretrained("gpt2")
model = PeftModel.from_pretrained(base_model, "isashap/contexttrained-validationloss-gpt2FINAL")

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# Define the prediction function
def predict_stock(input_text):
    # Tokenize the input
    inputs = tokenizer(input_text, return_tensors="pt")
    # Generate output using the model
    outputs = model.generate(**inputs, max_length=50, num_return_sequences=1)
    # Decode the output
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

demo = gr.Interface(
    fn=predict_stock,
    inputs=gr.Textbox(lines=2, label="Input Text (e.g., stock-related news)"),
    outputs=gr.Textbox(label="Prediction"),
    title="Stock Prediction with PEFT Model",
    description="A demo that uses a fine-tuned GPT-2 model with PEFT to predict stock outcomes based on input text."
    )    


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