Ai_project / app.py
Oreo99's picture
Create app.py (#1)
b24b7bc verified
import gradio as gr
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7")
model = AutoModelForCausalLM.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7")
# Define function to generate response
def generate_response(message, history, system_prompt, tokens):
# Concatenate system prompt and user message
input_text = f"{system_prompt} {message}"
# Tokenize input text
input_ids = tokenizer.encode(input_text, return_tensors="pt")
# Generate response
output = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
response = tokenizer.decode(output[0], skip_special_tokens=True)
return response
# Define Gradio interface
with gr.Blocks() as demo:
system_prompt = gr.Textbox("You are helpful AI.", label="System Prompt")
slider = gr.Slider(10, 100, render=False, label="Number of Tokens")
gr.ChatInterface(
generate_response,
inputs=["text", "text", system_prompt, slider],
outputs="text"
)
# Launch Gradio interface
demo.launch()