app.py

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, pipeline
import gradio as gr
from huggingface_hub import login
import os

# Huggingface Authentication
api_key = os.getenv('cs549_naacp_access_key')

if api_key is None:
    raise ValueError("API_KEY is not set in the environment variables.")

login(api_key)

# Load the model
model = AutoPeftModelForCausalLM.from_pretrained("Moritz-Pfeifer/financial-times-classification-llama-2-7b-v1.3")
tokenizer = AutoTokenizer.from_pretrained("Moritz-Pfeifer/financial-times-classification-llama-2-7b-v1.3")

def predict_text(test, model, tokenizer):
    prompt = f"""
            You are given a news article regarding the greater Boston area.
            Analyze the sentiment of the article enclosed in square brackets,
            determine if it is positive, negative or other, and return the answer as the corresponding sentiment label
            "positive" or "negative". If the sentiment is neither positive or negative, return "other".

            [{test}] ="""
    pipe = pipeline(task="text-generation",
                        model=model,
                        tokenizer=tokenizer,
                        max_new_tokens = 1,
                        temperature = 0.1,
                       )
    result = pipe(prompt)
    answer = result[0]['generated_text'].split("=")[-1]
    # print(answer)
    if "positive" in answer.lower():
        return "positive"
    elif "negative" in answer.lower():
        return "negative"
    else:
        return "other"

def predict(input_text):
  return predict_text(input_text, model, tokenizer)


interface = gr.Interface(fn=predict, inputs="text", outputs="text", title="Text Classifier", description="Insert your text and get the classification result.")
interface.launch()

if __name__ == "__main__":
    interface.launch(share=True)