| import gradio as gr |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| import torch |
|
|
| model_name = "your-username/your-model-name" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained(model_name) |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model.to(device) |
|
|
| def chat_fn(user_input, history=[]): |
| |
| conversation = "" |
| for u, r in history: |
| conversation += f"User: {u}\nBot: {r}\n" |
| conversation += f"User: {user_input}\nBot:" |
|
|
| inputs = tokenizer(conversation, return_tensors="pt").to(device) |
| outputs = model.generate(**inputs, max_length=500, pad_token_id=tokenizer.eos_token_id) |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Bot:")[-1].strip() |
|
|
| history.append((user_input, response)) |
| return history, history |
|
|
| iface = gr.ChatInterface(fn=chat_fn, title="Awesome's Chatbot") |
| iface.launch() |
|
|