AI_Discord_Bot / app.py
abcsnoob's picture
Create app.py
1b50568 verified
raw
history blame
613 Bytes
from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
import gradio as gr
model_id = "TinyLlama/TinyLlama-1.1B-Chat-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer, max_new_tokens=200, do_sample=True)
def chat(user_input):
res = pipe(user_input)[0]["generated_text"]
return res[len(user_input):] # chỉ trả lời, không lặp lại câu hỏi
gr.Interface(fn=chat, inputs="text", outputs="text", title="🤖 TinyLlama Chatbot").launch()