loopnews9/chat-bot

1import gradio as gr
2from transformers import AutoModelForCausalLM, AutoTokenizer
3import torch
4
5# Load Hugging Face model
6model_name = "meta-llama/Llama-2-7b-chat-hf"
7tokenizer = AutoTokenizer.from_pretrained(model_name)
8model = AutoModelForCausalLM.from_pretrained(
9    model_name, 
10    device_map="auto", 
11    torch_dtype=torch.float16
12)
13
14# Define chatbot logic
15def chatbot_response(user_message):
16    inputs = tokenizer(user_message, return_tensors="pt").to("cuda")
17    outputs = model.generate(inputs.input_ids, max_length=150, temperature=0.8, top_p=0.95)
18    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
19    return response
20
21# Gradio UI
22iface = gr.Interface(
23    fn=chatbot_response,
24    inputs=gr.Textbox(lines=2, placeholder="Type your message..."),
25    outputs="text",
26    title="Romantic Chatbot"
27)
28
29if __name__ == "__main__":
30    iface.launch()
31