retrAIced / pages /Text Generation.py
JavierGon12's picture
Text Genereration script for chatbot
96c3959
raw
history blame
1.39 kB
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)