Update app.py

app.py

@@ -1,20 +1,24 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 import torch
 
-model_id = "olacode55/zimble-llama2"
+base_model = "meta-llama/Llama-2-7b-chat-hf"
+adapter_model = "olacode55/zimble-llama2"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float16,
+tokenizer = AutoTokenizer.from_pretrained(base_model)
+base = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     device_map="auto"
 )
 
+model = PeftModel.from_pretrained(base, adapter_model)
+
 def generate(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7, do_sample=True)
+    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-demo = gr.Interface(fn=generate, inputs="text", outputs="text", title="Zimble LLaMA 2 Text Generator")
+demo = gr.Interface(fn=generate, inputs="text", outputs="text", title="Zimble LLaMA 2 Fine-Tuned")
 demo.launch()