Update app.py

app.py

@@ -3,5 +3,42 @@ import gradio as gr
 def greet(name):
     return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM,  BitsAndBytesConfig
+from peft import PeftModel, PeftConfig
+
+class InferenceFineTunning:
+    def __init__(self, model_path):
+        peft_model_id = f"hyang0503/{model_path}"
+        config = PeftConfig.from_pretrained(peft_model_id)
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16
+        )
+        self.model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, quantization_config=bnb_config, device_map="auto")
+        self.model = PeftModel.from_pretrained(self.model, peft_model_id)
+        
+        # self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+        self.tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.model.eval()
+
+    def generate(self, q): # 실습 노트북과 내용 다름
+        outputs = self.model.generate(
+            **self.tokenizer(
+                f"### 질문: {q}\n\n### 답변:",
+                return_tensors='pt',
+                return_token_type_ids=False
+            ).to("cuda"),
+            max_new_tokens=256,
+            early_stopping=True,
+            do_sample=True,
+            eos_token_id=2,
+        )
+        print(self.tokenizer.decode(outputs[0]))
+ifg = InferenceFineTunning("qlora-koalpaca")
+iface = gr.Interface(fn=ifg.generate, inputs="text", outputs="text")
 iface.launch()