Update app.py

app.py

@@ -1,32 +1,64 @@
-import gradio as gr
-import torch
-from gpt_dev import BigramLanguageModel  # Import your model class
-
-# Load model
-model = BigramLanguageModel()
-model.load_state_dict(torch.load('transformer_weights.pth'))
-model.eval()
-
-def predict(input_text):
-    # Your preprocessing and prediction code
-    # output = model(input_text)
-    tokens = tokenize(input_text)  # Define this function based on your tokenizer
-    input_tensor = torch.tensor(tokens).unsqueeze(0)  # Add batch dimension
-
-    with torch.no_grad():
-        output = model(input_tensor)
-    
-    # Postprocess output
-    predicted_tokens = torch.argmax(output, dim=-1).squeeze().tolist()
-    predicted_text = detokenize(predicted_tokens)  # Define this function based on your tokenizer
-    
-    return predicted_text  # Replace with actual prediction logic
-def tokenize(text):
-    # Example tokenizer function
-    return [ord(char) for char in text]  # Replace with actual tokenizer
-
-def detokenize(tokens):
-    # Example detokenizer function
-    return ''.join([chr(token) for token in tokens])  # Replace with actual detokenizer
-interface = gr.Interface(fn=predict, inputs="text", outputs="text")
-interface.launch()
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import PreTrainedModel, AutoConfig
+from pathlib import Path
+
+
+def predict_with_custom_weights(model_name, text, weights_path):
+  """
+  Loads a model config, creates a model, loads custom weights, performs testing 
+  and prediction.
+
+  Args:
+      model_name: Name of the pre-trained model architecture (e.g., "bert-base-uncased").
+      text: The text string for prediction.
+      weights_path: Path to the directory containing your custom weights.
+
+  Returns:
+      A dictionary containing predictions and logits (optional).
+  """
+  # Load tokenizer and model config
+  tokenizer = AutoTokenizer.from_pretrained(model_name)
+  config = AutoConfig.from_pretrained(model_name)
+
+  # Create empty model from config
+  model = PreTrainedModel.from_config(config)
+
+  # Load weights from your directory
+  model.load_state_dict(torch.load(Path(weights_path) / "pytorch_model.bin"))
+
+  # Tokenize the text
+  inputs = tokenizer(text, padding="max_length", truncation=True, return_tensors="pt")
+
+  # Set model to evaluation mode
+  model.eval()
+
+  # Make predictions
+  with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predictions = torch.argmax(logits, dim=-1)
+
+  # Get label names (optional, if your model has labels)
+  label_names = config.label_names if hasattr(config, "label_names") else None
+
+  # Return results
+  return {
+      "predictions": predictions.item() if len(predictions) == 1 else predictions.tolist(),
+      "logits": logits.squeeze().tolist() if label_names else None,
+      "labels": label_names[predictions.item()] if label_names else None,
+  }
+
+
+# Example usage (replace with your actual model name, text, and weights path)
+model_name = "bert-base-uncased"
+text = "This movie was absolutely fantastic!"
+weights_path = "transformer_weights.pth"  # Replace with actual path
+results = predict_with_custom_weights(model_name, text, weights_path)
+
+print(f"Predicted label: {results['predictions']}")
+
+if results.get("labels"):
+  print(f"Corresponding label name: {results['labels']}")
+
+if results.get("logits"):
+  print(f"Logits: {results['logits']}")