Update app.py

app.py

@@ -1,63 +1,10 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForPreTraining, pipeline
-from huggingface_hub import InferenceClient
-from typing import List, Tuple
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForPreTraining, AutoTokenizer
 
-# Load the InLegalBERT model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("law-ai/InLegalBERT")
-model = AutoModelForPreTraining.from_pretrained("law-ai/InLegalBERT")
+# Download model and tokenizer
+model_path = hf_hub_download("law-ai/InLegalBERT", filename="pytorch_model.bin")
+tokenizer_path = hf_hub_download("law-ai/InLegalBERT", filename="tokenizer_config.json")
 
-# Initialize the Gradio interface
-def respond(
-    message: str,
-    history: List[Tuple[str, str]],  # Using List and Tuple for type annotation
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    # Start with system message
-    messages = [{"role": "system", "content": system_message}]
-    
-    # Append history to the messages list
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    # Append the current user message
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    # Use the tokenizer to process the input message
-    inputs = tokenizer(message, return_tensors="pt")
-
-    # Use the model to generate a response
-    with torch.no_grad():  # Disable gradients since we're just doing inference
-        outputs = model(**inputs)
-        response = tokenizer.decode(outputs.logits.argmax(dim=-1)[0], skip_special_tokens=True)
-
-    yield response
-
-# Set up Gradio interface with additional options
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-# Launch the Gradio app
-if __name__ == "__main__":
-    demo.launch()
+# Load locally downloaded model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+model = AutoModelForPreTraining.from_pretrained(model_path)