import streamlit as st
from PIL import Image
import streamlit as st
from transformers import pipeline
import pandas as pd
import plotly.express as px
import matplotlib.pyplot as plt
from pathlib import Path
import base64
from st_pages import Page, add_page_title, show_pages
from streamlit_extras.badges import badge
import transformers




model_name = 'Intel/neural-chat-7b-v3-1'
model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

def generate_response(system_input, user_input):

    # Format the input using the provided template
    prompt = f"### System:\n{system_input}\n### User:\n{user_input}\n### Assistant:\n"

    # Tokenize and encode the prompt
    inputs = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False)

    # Generate a response
    outputs = model.generate(inputs, max_length=1000, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract only the assistant's response
    return response.split("### Assistant:\n")[-1]


# Example usage
system_input = "You are a employee in the customer succes department of a company called Retraced that works in sustainability and traceability"
prompt = st.text_input(str("Insert here you prompt?"))
response = generate_response(system_input, prompt)
st.write(response)