Update app.py

app.py

@@ -1,14 +1,32 @@
 import torch
+from transformers import pipeline, logging, AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
-from utils import *
-from torch import nn
-import lightning.pytorch as pl
-from torch.nn import functional as F
 
-device     = 'cuda' if torch.cuda.is_available() else 'cpu'
+## 1 - Loading Model
+model_name = "microsoft/phi-2"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    device_map='auto',
+)
+model.config.use_cache = False
 
-HTML_TEMPLATE = """    
-<style>
+## 2 - Loading Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
+
+## 3 - Load adapter (trained LORA weights)
+peft_model_folder = 'checkpoint700'
+model.load_adapter(peft_model_folder)
+
+def generate_dialogue(input_text):
+
+  pipe = pipeline(task="text-generation",model=model,tokenizer=tokenizer,max_length=200)
+  result = pipe(f"<s>[INST] {input_text} [/INST]")
+  return result[0]['generated_text']
+    
+HTML_TEMPLATE = """
+  <style>
     
     #app-header {
         text-align: center;
@@ -60,9 +78,9 @@ HTML_TEMPLATE = """
     <div class="artifact large"></div>
     <div class="artifact large"></div>
     <!-- Content -->
-    <h1>GPT NEXT WORD GENERATOR</h1>
+    <h1>CHAT with fine tuned Phi-2 LLM</h1>
     <p>Generate dialogue for given some initial prompt for context.</p>
-    <p>Model: GPT, Dataset: arxiv + book + cc, Parameter Count: 160M</p>
+    <p>Model: Phi-2 (https://huggingface.co/microsoft/phi-2),  Dataset: oasst1 (https://huggingface.co/datasets/OpenAssistant/oasst1) </p>
 """
 
 with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background: url('file=https://github.com/santule/ERA/assets/20509836/e78f2bb3-ddd8-4ce9-a941-3d3d7ef7a272')}") as interface:
@@ -89,22 +107,13 @@ with gr.Blocks(theme=gr.themes.Glass(),css=".gradio-container {background: url('
             value="Enter your prompt here: This text will set the context for the AI's response."
         )
 
-        temperature_dropdown = gr.Slider(0, 1, value=0.8, label="Temperature", info="Set the creativity level: Higher values produce more varied results, lower values generate more predictable text.")
-        top_k_dropdown = gr.Slider(200, 300, value=200, label="Top K", info="Control the randomness: Limits the AI to consider only the top K most likely next words.")
-        max_new_tokens = gr.Slider(10, 100, value=50, label="Max Tokens", info="Choose the length: This determines the maximum number of words the AI will generate.")
-
-
         outputs = gr.Textbox(
-            label="Generated Dialogue"
+            label="Answer"
         )
-        inputs = [input_text, temperature_dropdown, top_k_dropdown, max_new_tokens]
+        inputs = [input_text]
    
     with gr.Column():
-        button = gr.Button("Generate")
+        button = gr.Button("Ask me")
         button.click(generate_dialogue, inputs=inputs, outputs=outputs)
 
-    # with gr.Row():
-    #     gr.Examples(examples=examples, inputs=inputs, outputs=outputs, fn=generate_dialogue, cache_examples=True,)
-
-
 interface.launch()