app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Loading Llama model
model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map="auto", torch_dtype=torch.float16
)

# Answer generation
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(inputs.input_ids, max_length=200, temperature=0.7)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio interface
demo = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="Llama Chatbot",
    description="Chatbot Llama-2-7b-hf",
)

if __name__ == "__main__":
    demo.launch()