Lukeam/aa_book

1import os2import gradio as gr3from transformers import AutoModelForCausalLM, AutoTokenizer4from peft import PeftModel5from huggingface_hub import login6 7def generate_response(question):8    try:9        # Debug print to verify token10        token = os.environ.get('HUGGINGFACE_TOKEN')11        print(f"Token available: {'Yes' if token else 'No'}")12        login(token)13        14        print("Loading tokenizer...")15        tokenizer = AutoTokenizer.from_pretrained(16            "meta-llama/Llama-2-7b-hf",17            use_auth_token=True  # Changed to use environment token18        )19        20        print("Loading model...")21        base_model = AutoModelForCausalLM.from_pretrained(22            "meta-llama/Llama-2-7b-hf",23            use_auth_token=True,  # Changed to use environment token24            load_in_8bit=True,25            device_map="auto"26        )27        model = PeftModel.from_pretrained(base_model, "Lukeam/llama-aa-fine-tuned")28        29        # Format prompt30        prompt = f"### Question: {question}\n\n### Answer:"31        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")32        33        # Generate response34        outputs = model.generate(35            **inputs,36            max_length=512,37            temperature=0.7,38            num_return_sequences=139        )40        41        response = tokenizer.decode(outputs[0], skip_special_tokens=True)42        return response.split("### Answer:")[1].strip()43        44    except Exception as e:45        return f"Error generating response: {str(e)}"