Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
-import time
-import torch
-
+from gradio.components import Textbox
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GenerationConfig
+
 
 peft_model_id = "Ngadou/falcon-7b-scam-buster"
 config = PeftConfig.from_pretrained(peft_model_id)
@@ -11,37 +11,8 @@ config = PeftConfig.from_pretrained(peft_model_id)
 model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True, return_dict=True, load_in_4bit=True, device_map='auto')
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
+# Adapter model
 model = PeftModel.from_pretrained(model, peft_model_id).to("cuda")
-
-
-
-# Load the Lora model
-model = PeftModel.from_pretrained(model, peft_model_id)
-
-tokenizer.pad_token = tokenizer.eos_token
-
-def generate(chat):
-
-    input_text = chat + "\nIs this conversation a scam or not and why?"
-    
-    
-    encoding = tokenizer(input_text, return_tensors="pt").to("cuda")
-    output = model.generate(
-      input_ids=encoding.input_ids, 
-      attention_mask=encoding.attention_mask, 
-      max_new_tokens=100, 
-      do_sample=True, 
-      temperature=0.000001, 
-      eos_token_id=tokenizer.eos_token_id, 
-      top_k = 0
-      )
-    
-    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    output_text = output_text.replace(example_text, "").lstrip("\n")
-    
-    print("\nAnswer:")
-    print(output_text)
-    return output_text
   
 
 # def is_scam(instruction):
@@ -83,11 +54,45 @@ def generate(chat):
 #     return classification #, reason
 
 
-# Define the Gradio interface
+def is_scam(instruction):
+    max_new_tokens=128
+    temperature=0.1
+    top_p=0.75
+    top_k=40
+    num_beams=4
+
+    instruction = instruction + ".\nIs this conversation a scam or not and why?"
+    prompt = instruction + "\n### Solution:\n"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to("cuda")
+    attention_mask = inputs["attention_mask"].to("cuda")
+    generation_config = GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        num_beams=num_beams,
+    )
+    with torch.no_grad():
+        generation_output = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            output_scores=True,
+            max_new_tokens=max_new_tokens,
+            early_stopping=True
+        )
+    s = generation_output.sequences[0]
+    output = tokenizer.decode(s)
+
+    classification = output.split("### Solution:")[1].lstrip("\n")
+    print(classification)
+
+    return str(classification), "Hello World"
+
+
 gr.Interface(
-    fn=generate,
+    fn=is_scam,
     inputs='text',
-    outputs=[
-        gr.outputs.Textbox(label="Classification and rational")
-    ]
-).launch()
+    outputs= ['text','text']
+).launch(share=True, debug=True)