File size: 1,897 Bytes
c4868a7 d91393a fe52cbc d91393a c4868a7 d91393a f192afa d91393a c4868a7 d91393a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load the LLaMA-2 model and tokenizer from Hugging Face
model_name = "meta-llama/Llama-2-7b-hf" # Change to the desired LLaMA model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
# Function to generate responses
def generate_response(user_input, chat_history):
# Add the user's input to the conversation history
chat_history.append({"role": "user", "content": user_input})
# Prepare input for the model
conversation = ""
for turn in chat_history:
conversation += f"{turn['role']}: {turn['content']}\n"
inputs = tokenizer(conversation, return_tensors="pt").to(model.device)
# Generate model response
outputs = model.generate(inputs.input_ids, max_length=500, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Add the model's response to the chat history
chat_history.append({"role": "assistant", "content": response})
# Only return the model's response for display
return response, chat_history
# Initialize the chat history
chat_history = []
# Define Gradio interface
with gr.Blocks() as chat_interface:
gr.Markdown("## LLaMA-2 Chatbot")
chat_input = gr.Textbox(label="Your Message")
chat_output = gr.Chatbot()
# Update chat on button click
def handle_input(user_input):
response, chat_history = generate_response(user_input, chat_history)
chat_output.update(chat_history)
return "", chat_history # Clear input box and update chat history
chat_input.submit(handle_input, inputs=chat_input, outputs=[chat_input, chat_output])
# Launch Gradio app
chat_interface.launch()
|