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

model_name = 'ch103/llama-2-7b-miniguanaco'
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name)

pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)

def generate_response(user_input):
    formatted_input = f"<s>[INST] {user_input} [/INST]"
    result = pipe(formatted_input)
    response_text = result[0]['generated_text']

    # Extracting the AI response part, assuming it follows a specific format after [/INST]
    response_parts = response_text.split("[/INST]")
    if len(response_parts) > 1:
        ai_response = response_parts[1].split("</s>")[0].strip()  # Assuming the response ends with </s>
    else:
        ai_response = "Sorry, I couldn't process that message."

    return ai_response

gradio_app = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
    outputs=gr.Textbox(),
    title="Text-Based Interaction Interface",
    description="Interact using text. Type your input and receive a response."
)

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