Modified for unsloth

app.py

@@ -2,20 +2,42 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import nltk
 nltk.download('punkt')
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer
 
-def generate_subject(model_name,question_body):
-  tokenizer = AutoTokenizer.from_pretrained(model_name)
-  model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-  inputs = ["Ask a question: " + question_body]
-  inputs = tokenizer(inputs, max_length=512, truncation=True, return_tensors="pt")
-  output = model.generate(**inputs, num_beams=8, do_sample=True, min_length=1, max_length=10)
-  decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
-  predicted_title = nltk.sent_tokenize(decoded_output.strip())[0]
+def preprocess_text(text):
+    # Convert to lowercase
+    text = text.lower()
+    # Remove punctuation
+    text = re.sub(r'[^\w\s]', '', text)
+    # Remove extra whitespace
+    text = ' '.join(text.split())
+    return text
+
+def generate_answer(model_name,question):
+    model = AutoPeftModelForCausalLM.from_pretrained(
+        model_name, # YOUR MODEL YOU USED FOR TRAINING
+        load_in_4bit = load_in_4bit,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    question_preprocessed = preprocess_text(question)
+    inputs = tokenizer(
+    [
+        qa_prompt.format(
+        "Please provide the answer for the question", # instruction
+        question_preprocessed, # input
+        "", # output - leave this blank for generation!
+        )
+    ], return_tensors = "pt")
+    outputs = model.generate(**inputs, max_new_tokens = 64, use_cache = True)
+    decoded_output = tokenizer.batch_decode(outputs,skip_special_tokens=True)[0]
+    predicted_title = nltk.sent_tokenize(decoded_output.strip())[0]
   return predicted_title
 
 
 iface = gr.Interface(
-    fn=generate_subject,
+    fn=generate_answer,
     inputs=[
         gr.Dropdown(choices=["GSridhar1982/AIML_QA_Mistral7B_FineTuned_Unsloth","GSridhar1982/AIML_QA_Mistral7B_FineTuned_Unsloth"], label="Select Model"),
         gr.Textbox(lines=5, label="Question")