Update app.py

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import torch
-from peft import PeftModel, PeftConfig
-from transformers import AutoTokenizer
+from peft import PeftModel, PeftConfig, LoraConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import load_dataset
+from trl import SFTTrainer
 
 ref_model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m-deduped-v0", torch_dtype=torch.bfloat16)
 peft_model_id = "w601sxs/pythia-70m-instruct-orca-chkpt-64000"
@@ -13,12 +15,12 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 model.eval()
 
 def predict(text):
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
-    outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10)
-    out_text = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
+        outputs = model.generate(input_ids=inputs["input_ids"], max_new_tokens=10)
+        out_text = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]
     
-    return out_text
+    return out_text.split(text)[-1]
     
 
 demo = gr.Interface(
@@ -27,6 +29,4 @@ demo = gr.Interface(
     outputs='text',
 )
 
-demo.launch()
-
-
+demo.launch()