HackingBot / app.py
data-ject's picture
Update app.py
d926646 verified
raw
history blame
1.42 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Load model and tokenizer from Hugging Face Model Hub
model_name = "meta-llama/Meta-Llama-3.1-70B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Define system instruction
system_instruction = "You are a helpful assistant. Provide detailed and accurate responses to the user's queries."
# Define the chat function
def chat_function(prompt):
# Create the full input prompt including the system instruction
full_prompt = f"{system_instruction}\nUser: {prompt}\nAssistant:"
# Tokenize the full prompt
inputs = tokenizer(full_prompt, return_tensors="pt")
# Generate model response
with torch.no_grad():
outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
# Decode and return response
response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
# Extract only the assistant's response
response = response.split("Assistant:")[-1].strip()
return response
# Create Gradio interface
iface = gr.Interface(
fn=chat_function,
inputs="text",
outputs="text",
title="Meta-Llama Chatbot",
description="A chatbot powered by the Meta-Llama-3.1-70B-Instruct model."
)
# Launch the interface
if __name__ == "__main__":
iface.launch()