bug fix

app.py

@@ -1,24 +1,13 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline,BitsAndBytesConfig
-
-#model = AutoModelForCausalLM.from_pretrained("checkpoint_500",trust_remote_code=True)
-
-model_name = "microsoft/phi-2"
-
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import os
 token = os.environ.get("HUGGING_FACE_TOKEN")
 
-
-#bnb_config = BitsAndBytesConfig(
-#    load_in_4bit=True,
-#    bnb_4bit_quant_type="nf4",
-#    bnb_4bit_compute_dtype=torch.float16,
-#)
+model_name = "microsoft/phi-2"
 
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    #quantization_config=bnb_config,
     use_auth_token=token,
     trust_remote_code=True
 )
@@ -28,25 +17,28 @@ model.load_adapter("checkpoint_500")
 tokenizer = AutoTokenizer.from_pretrained("checkpoint_500", trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
+
 def inference(prompt, count):
     count = int(count)
     pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer) 
     result = pipe(f"{prompt}",max_new_tokens=count)
-    out_text = result[0]['generated_text']
-    return out_text
+    output = result[0]['generated_text']
+    return output
 
-title = "TSAI S21 Assignment: Adaptive QLoRA training on open assist oasst1 dataset, using microsoft/phi2 model"
-description = "A simple Gradio interface that accepts a context and generates GPT like text "
-examples = [["What is a large language model?","50"]
-           ]
- 
+examples = [
+    ["What is LLM?","50"]
+]
 
 demo = gr.Interface(
     inference, 
-    inputs = [gr.Textbox(placeholder="Enter a prompt"), gr.Textbox(placeholder="Enter number of characters you want to generate")], 
-    outputs = [gr.Textbox(label="Chat GPT like text")],
-    title = title,
-    description = description,
+    inputs = [
+        gr.Textbox(placeholder="Enter a prompt"), 
+        gr.Textbox(placeholder="Enter number of characters you want to generate")
+    ], 
+    outputs = [
+        gr.Textbox(label="Generated text")
+    ],
     examples = examples
 )
+
 demo.launch()