Create app.py

app.py

@@ -1,7 +1,39 @@
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
 
-def greet(name):
-    return "Hello " + name + "!!"
+base_model = "mistralai/Mistral-7B-v0.1"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit= True,
+    bnb_4bit_quant_type= "nf4",
+    bnb_4bit_compute_dtype= torch.bfloat16,
+    bnb_4bit_use_double_quant= False,
+)
+model = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    quantization_config=bnb_config,
+    device_map={"": 0}
+)
+
+ft_model = PeftModel.from_pretrained(model, 'kiki7sun/mixtral-academic-finetune-QLoRA-0121')
+
+tokenizer = AutoTokenizer.from_pretrained(
+    base_model_id,
+    add_bos_token=True,
+    trust_remote_code=True,
+)
+
+ft_model.eval()
+def greet(eval_prompt,max_new_tokens):
+  model_input = tokenizer(eval_prompt, return_tensors="pt").to("cuda")
+  with torch.no_grad():
+      generation = ft_model.generate(**model_input, max_new_tokens = max_new_tokens)
+      result = tokenizer.decode(generation[0], skip_special_tokens=True)
+
+  return result
+
+demo = gr.Interface(fn=greet, inputs="textbox", outputs="textbox")
+
+demo.queue().launch(debug=True, share=True, inline=False)