from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import gradio as gr
import spaces

# --- Model Loading ---
base_model_id = "unsloth/Meta-Llama-3.1-8B"
lora_model_id = "Nlpeva/lora_model"  # Replace with your LoRA Hub path

try:
    model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
    model = PeftModel.from_pretrained(model, lora_model_id)
    print("Model and LoRA loaded successfully!")
except Exception as e:
    print(f"Error loading model or LoRA: {e}")
    model = None
    tokenizer = None

# --- Generation Function ---
@spaces.GPU
def generate_response(information, input_text):
    if model is None or tokenizer is None:
        return "Model not loaded. Please check the logs."

    prompt = f"Information: {information}\n\nInput: {input_text}\n\nResponse:"
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device)

    try:
        with torch.no_grad():
            output = model.generate(
                input_ids=input_ids,
                max_length=300,  # Adjust as needed
                num_return_sequences=1,
                temperature=0.7,
                top_p=0.9,
                # Add other generation parameters as desired
            )
        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        return generated_text.strip()
    except Exception as e:
        return f"Error during generation: {e}"

# --- Gradio Interface ---
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Information", placeholder="Provide any relevant context or information here."),
        gr.Textbox(label="Input", placeholder="Enter your query or the text you want the model to process.")
    ],
    outputs=gr.Textbox(label="Output"),
    title="Llama-3 with Custom LoRA",
    description="Enter information and an input, and the model will generate a response based on both."
)

iface.launch()