app.py

import gradio as gr
from transformers import LlamaTokenizer, LlamaForCausalLM, pipeline

# Load your model and tokenizer
model_name = "midrees2806/Llama-2-7b-updatedchatbot-finetune"
tokenizer = LlamaTokenizer.from_pretrained(model_name)
model = LlamaForCausalLM.from_pretrained(model_name)

# Define the pipeline
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer)

# Define the function to generate responses
def generate_response(prompt):
    # Format the prompt as required by the model
    input_text = f"<s>[INST] {prompt} [/INST]"
    response = pipe(input_text)
    # Extract the generated text from the response
    answer = response[0]['generated_text'].split('[/INST]')[-1].strip()
    return answer

# Gradio Interface setup
iface = gr.Interface(
    fn=generate_response,
    inputs="text",
    outputs="text",
    title="LLaMA-2 Chatbot",
    description="Ask anything to the LLaMA-2 fine-tuned model!",
)

# Launch the Gradio app
iface.launch()