Update app to fine-tuned Llama chat interface

app.py

@@ -1,7 +1,81 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
+import os
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Configuration
+MODEL_ID = "meta-llama/Llama-3.2-1B"
+ADAPTER_ID = "Zedel17/fine_tuned_llama_1b"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# Load model with 4-bit quantization
+print("Loading model...")
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+# Get HF token from Space secrets
+hf_token = os.environ.get("HF_TOKEN")
+
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=bnb_config,
+    device_map="auto",
+    token=hf_token
+)
+
+# Load fine-tuned adapter
+model = PeftModel.from_pretrained(base_model, ADAPTER_ID, token=hf_token)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=hf_token)
+tokenizer.pad_token = tokenizer.eos_token
+
+print("Model loaded successfully!")
+
+def chat(message, history):
+    # Format the conversation
+    conversation = ""
+    for user_msg, bot_msg in history:
+        conversation += f"human: {user_msg}\ngpt: {bot_msg}\n"
+    conversation += f"human: {message}\ngpt: "
+
+    # Tokenize and generate
+    inputs = tokenizer(conversation, return_tensors="pt", truncation=True, max_length=512).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract just the new response
+    response = response.split("gpt: ")[-1].strip()
+
+    return response
+
+# Create Gradio interface
+demo = gr.ChatInterface(
+    fn=chat,
+    title="Fine-Tuned Llama 1B Chat",
+    description="Chat with a fine-tuned Llama-3.2-1B model",
+    examples=[
+        "Hello! How are you?",
+        "Can you help me with a coding question?",
+        "Tell me something interesting",
+    ],
+    retry_btn=None,
+    undo_btn="Delete Previous",
+    clear_btn="Clear",
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,8 @@
+gradio
+transformers>=4.35.0
+torch>=2.0.0
+accelerate
+peft
+bitsandbytes
+sentencepiece
+protobuf