|
import gradio as gr |
|
import time |
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7") |
|
model = AutoModelForCausalLM.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7") |
|
|
|
|
|
def generate_response(message, history, system_prompt, tokens): |
|
|
|
input_text = f"{system_prompt} {message}" |
|
|
|
input_ids = tokenizer.encode(input_text, return_tensors="pt") |
|
|
|
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 |
|
|
|
|
|
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" |
|
) |
|
|
|
|
|
demo.launch() |
|
|